From 4204a9e48237a96496fb76bed890ef77dd464923 Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Wed, 21 Oct 2026 08:07:38 +0000 Subject: [PATCH 01/18] fix(offload): make the state-tier CPU budget a reservation in the shared LMCache pool StateByteCodec bounds state images with OFFLOAD_STATE_CPU_SIZE, but the images live in the same LocalCPU hot cache as paged KV and LMCache's LRU can still evict them under KV pressure. StateOffloadIndex keeps advertising those hashes, so the next request over that prefix parks a joint load whose state leg misses, and the whole prompt is recomputed. Under a budget, hold one extra MemoryObj reference per stored image. LocalCPU only evicts objects whose ref count is 1, so only the codec's own LRU retires an image (remove, then drop the reference). A put that LocalCPU refuses (key already resident, hot cache off) keeps no reference, and a re-store of a resident image is skipped. Co-authored-by: Cursor --- .../offload/hybrid/kimi_k3/state_object.py | 47 ++++- tests/test_state_tier_reservation.py | 190 ++++++++++++++++++ 2 files changed, 232 insertions(+), 5 deletions(-) create mode 100644 tests/test_state_tier_reservation.py diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py index 18094ed792..c9e19905b8 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py @@ -66,6 +66,12 @@ def __init__( if self._max_entries is not None and self._max_entries < 1: raise ValueError("state CPU budget must fit at least one checkpoint") self._cache_keys: OrderedDict[Any, None] = OrderedDict() + # Under a budget, one extra reference per stored image. LMCache's + # LocalCPU LRU only evicts objects whose ref count is 1, so the pool + # cannot drop an image behind the codec's back while the engine-side + # `StateOffloadIndex` still advertises it. The budget is a reservation: + # only `put`'s own LRU retires these, KV pressure cannot. + self._held: dict[Any, Any] = {} self._cache_lock = threading.Lock() self._store_lock = threading.Lock() # Never hard-code a size: V4 keeps six compressor fields across @@ -126,11 +132,27 @@ def put(self, h: int, unit_ids, on_source_released=None) -> bool: key = self.key(h) with self._store_lock: with self._cache_lock: - if key not in self._cache_keys: - while len(self._cache_keys) >= self._max_entries: - old = next(iter(self._cache_keys)) - self._storage.remove(old) - del self._cache_keys[old] + if key in self._held: + # Same prefix hash and layout, same bytes, and the held + # reference keeps the stored copy resident: storing again + # would only cost an image-sized D2H that LocalCPU then + # refuses as a duplicate key. + self._cache_keys.move_to_end(key) + stored = True + else: + stored = False + if key not in self._cache_keys: + while len(self._cache_keys) >= self._max_entries: + old = next(iter(self._cache_keys)) + self._storage.remove(old) + del self._cache_keys[old] + held = self._held.pop(old, None) + if held is not None: + held.ref_count_down() + if stored: + if on_source_released is not None: + on_source_released() + return True ok = self._put(h, unit_ids, on_source_released) if ok: with self._cache_lock: @@ -164,6 +186,9 @@ def _put(self, h: int, unit_ids, on_source_released=None) -> bool: obj = self._allocate(self.entry_bytes) if obj is None: return False + hold = self._max_entries is not None + if hold: + obj.ref_count_up() # `batched_put` discharges the reference it is handed, but only in its # terminal tail loop (LMCache `storage_manager.py`, one # `ref_count_down()` per object), which runs *after* every step that can @@ -196,7 +221,19 @@ def _put(self, h: int, unit_ids, on_source_released=None) -> bool: self._storage.batched_put([key], [obj]) except Exception: obj.ref_count_down() + if hold: + obj.ref_count_down() raise + if hold: + ref_count = getattr(obj, "get_ref_count", None) + if not callable(ref_count) or ref_count() > 1: + # The hot cache adopted it (its own reference plus ours). + with self._cache_lock: + self._held[key] = obj + else: + # Refused, e.g. the key was already resident: ours is the + # last reference, and keeping it would strand the block. + obj.ref_count_down() return True def get(self, h: int, slot) -> bool: diff --git a/tests/test_state_tier_reservation.py b/tests/test_state_tier_reservation.py new file mode 100644 index 0000000000..15983b6ac9 --- /dev/null +++ b/tests/test_state_tier_reservation.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: MIT +"""The state tier's CPU budget is a reservation inside the shared LMCache pool. + +State images and paged KV share one LocalCPU pool. Before, the codec's budget +only bounded the codec's own LRU: KV writes could still push LMCache's LRU over +the images, while the engine-side `StateOffloadIndex` kept advertising them, and +every such hit parked a joint load that then failed and recomputed the whole +prompt. These tests model LocalCPU's eviction rule (only objects whose ref count +is 1 are evictable) and check that a budgeted image survives KV pressure until +the codec itself retires it. +""" + +import dataclasses +import sys +import types +from types import SimpleNamespace + +import pytest + + +@pytest.fixture +def lmcache_key(monkeypatch): + @dataclasses.dataclass(frozen=True) + class CacheEngineKey: + model_name: str + world_size: int + worker_id: int + chunk_hash: int + dtype: object + + utils_module = types.ModuleType("lmcache.utils") + utils_module.CacheEngineKey = CacheEngineKey + lmcache_module = types.ModuleType("lmcache") + lmcache_module.__path__ = [] + lmcache_module.utils = utils_module + monkeypatch.setitem(sys.modules, "lmcache", lmcache_module) + monkeypatch.setitem(sys.modules, "lmcache.utils", utils_module) + return CacheEngineKey + + +class _Obj: + """A MemoryObj's reference count, and nothing else.""" + + def __init__(self, size: int = 1024) -> None: + self.refs = 1 + self.size = size + + def ref_count_up(self) -> None: + self.refs += 1 + + def ref_count_down(self) -> None: + assert self.refs > 0, "reference discharged twice" + self.refs -= 1 + + def get_ref_count(self) -> int: + return self.refs + + def get_size(self) -> int: + return self.size + + +class _LocalCPU: + """`storage_manager` over LocalCPU's rules. + + `batched_put` adopts an object with its own reference unless the key is + already resident, then drops the caller's reference. `evict_lru` is the + pool making room for KV: it skips anything another holder still references. + """ + + def __init__(self, refuse_all: bool = False) -> None: + self.hot: dict = {} + self.refuse_all = refuse_all + self.fail_put = False + + def batched_put(self, keys, objs) -> None: + if self.fail_put: + raise RuntimeError("put failed") + for k, o in zip(keys, objs, strict=True): + if not self.refuse_all and k not in self.hot: + o.ref_count_up() + self.hot[k] = o + o.ref_count_down() + + def get(self, key): + o = self.hot.get(key) + if o is not None: + o.ref_count_up() + return o + + def remove(self, key) -> int: + o = self.hot.pop(key, None) + if o is None: + return 0 + o.ref_count_down() + return 1 + + def evict_lru(self) -> list: + gone = [k for k, o in self.hot.items() if o.get_ref_count() == 1] + for k in gone: + self.hot.pop(k).ref_count_down() + return gone + + +def _codec(storage, max_entries): + from atom.kv_transfer.offload.hybrid.kimi_k3.state_object import StateByteCodec + + codec = StateByteCodec( + backend=SimpleNamespace( + page_unit_views=lambda units: [], state_entry_views=lambda slot: [] + ), + staged=SimpleNamespace(pack=lambda views, obj: None, unpack=lambda o, v: None), + entry_bytes=1024, + model_name="qwen", + world_size=1, + worker_id=0, + layout_id="L", + max_cache_bytes=None if max_entries is None else max_entries * 1024, + ) + codec.bind_storage_manager(storage) + allocated = [] + + def allocate(nbytes): + allocated.append(_Obj(nbytes)) + return allocated[-1] + + codec._allocate = allocate + return codec, allocated + + +def test_kv_pressure_cannot_evict_a_budgeted_image(lmcache_key): + pool = _LocalCPU() + codec, objs = _codec(pool, max_entries=2) + kv = _Obj() + pool.hot["kv-chunk"] = kv # a KV chunk only the pool references + + assert codec.put(1, []) + assert objs[0].refs == 2 # the pool's reference and the codec's + assert pool.evict_lru() == ["kv-chunk"] + assert codec.get(1, 0), "the image the index advertises must still load" + + +def test_the_codec_lru_retires_and_frees_its_oldest_image(lmcache_key): + pool = _LocalCPU() + codec, objs = _codec(pool, max_entries=2) + assert codec.put(1, []) and codec.put(2, []) and codec.put(3, []) + + assert codec.key(1) not in pool.hot + assert objs[0].refs == 0, "a retired image goes back to the allocator" + assert [o.refs for o in objs[1:]] == [2, 2] + assert len(codec._held) == 2 + + +def test_a_resident_image_is_not_stored_twice(lmcache_key): + pool = _LocalCPU() + codec, objs = _codec(pool, max_entries=2) + released = [] + assert codec.put(1, []) + assert codec.put(1, [], on_source_released=lambda: released.append(True)) + + assert len(objs) == 1, "no second image-sized allocation and D2H" + assert released == [True], "the source units must still be handed back" + + +def test_a_refused_put_keeps_no_reference(lmcache_key): + pool = _LocalCPU(refuse_all=True) + codec, objs = _codec(pool, max_entries=2) + assert codec.put(1, []) + + assert objs[0].refs == 0, "holding a refused object would strand its block" + assert codec._held == {} + + +def test_a_failed_put_discharges_both_references(lmcache_key): + pool = _LocalCPU() + pool.fail_put = True + codec, objs = _codec(pool, max_entries=2) + with pytest.raises(RuntimeError): + codec.put(1, []) + + assert objs[0].refs == 0 + assert codec._held == {} + + +def test_an_unbounded_codec_leaves_eviction_to_lmcache(lmcache_key): + pool = _LocalCPU() + codec, objs = _codec(pool, max_entries=None) + assert codec.put(1, []) + + assert objs[0].refs == 1 + assert pool.evict_lru() == [codec.key(1)] From 39dc5dda18b076f8cd9127035944b866ab70f769 Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Tue, 27 Oct 2026 08:00:32 +0000 Subject: [PATCH 02/18] fix(offload): refresh state images in LMCache's LRU on every tier hit Replace the fixed state-tier reservation with the fix the review asked for. LMCache's LocalCPU refreshes every KV chunk a lookup touches (contains(pin=True) -> touch_cache() -> cache_policy.update_on_hit), but a state hit is answered engine-side by StateOffloadIndex, so LMCache never saw an image being used. A prefix that hit every turn kept its KV at the young end and its image at the cold end, where pool pressure evicted it first, and the index kept advertising it until a joint load parked on it, missed and recomputed the whole prompt. Refresh: - StateOffloadIndex(refresh_hits=True) queues a touch for every hash could_serve votes a hit for (the one predicate every voter goes through), throttled to one per hash per 0.5 s. - The scheduler publishes them with the state loads; the K3 connector carries them as LMCacheOffloadMetadata.state_touches, each with a unique StateTouchId so a re-touch is a fresh aggregator key. - Each worker calls StateByteCodec.touch: update_on_hit under LocalCPU's cpu_lock (no pin, so nothing to unpin and no interference with the KV lookup's keys_in_request). StateByteCodec.get and a resident re-store refresh too; get_blocking alone does not. Index consistency: - Each rank reports resident/gone on the new k3_state_touch channel; failure-dominant across TP. A gone verdict forgets the hash. - record_allocation(defer_unverified=True) holds an admission that would load a CPU image not confirmed resident in the last 5 s; the scheduler returns it to waiting like an in-flight prefix. One pass later the verdict is in: a resident image loads, a gone one has left the index and the retry plans a lower boundary (HBM checkpoint, older image, or prefill) instead of parking on a miss. Bounded by OFFLOAD_LOOKUP_DEFER_S; a connector that cannot carry touches turns the refresh off. Budget: - No reservation: images hold no extra reference, and OFFLOAD_STATE_CPU_SIZE is only a cap whose codec LRU now follows use, not store order. - The worker warns at startup when the cap is not below the LMCache pool (it bounds nothing) or exceeds half of it. Tests: tests/test_state_tier_lru_refresh.py replaces the reservation tests. Co-authored-by: Cursor --- .../disaggregation/multi/multi_connector.py | 23 + atom/kv_transfer/disaggregation/types.py | 20 + atom/kv_transfer/offload/README.md | 14 + atom/kv_transfer/offload/connector.py | 13 + .../offload/hybrid/kimi_k3/connector.py | 97 ++- .../offload/hybrid/kimi_k3/state_object.py | 117 +-- atom/kv_transfer/offload/metadata.py | 6 + atom/model_engine/block_manager.py | 71 +- atom/model_engine/scheduler.py | 47 +- atom/model_engine/state_offload.py | 200 ++++- docs/environment_variables.md | 2 +- tests/test_lmcache_offload_connector.py | 3 + tests/test_state_tier_lru_refresh.py | 687 ++++++++++++++++++ tests/test_state_tier_reservation.py | 190 ----- 14 files changed, 1232 insertions(+), 258 deletions(-) create mode 100644 tests/test_state_tier_lru_refresh.py delete mode 100644 tests/test_state_tier_reservation.py diff --git a/atom/kv_transfer/disaggregation/multi/multi_connector.py b/atom/kv_transfer/disaggregation/multi/multi_connector.py index 3793c81082..dc9f233f69 100644 --- a/atom/kv_transfer/disaggregation/multi/multi_connector.py +++ b/atom/kv_transfer/disaggregation/multi/multi_connector.py @@ -644,6 +644,29 @@ 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_reports(self): """`(indexed, failed)` from the tier sub, else two empty sets. diff --git a/atom/kv_transfer/disaggregation/types.py b/atom/kv_transfer/disaggregation/types.py index 95f46a7b1d..9a45b0d4e7 100644 --- a/atom/kv_transfer/disaggregation/types.py +++ b/atom/kv_transfer/disaggregation/types.py @@ -121,12 +121,32 @@ 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, + because the aggregator tombstones each `(channel, operation_id)`. + """ + + 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] diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index 48cc1dbaf0..d74a624922 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -122,6 +122,20 @@ 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 admission about to load an image not confirmed in the last +few seconds 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. `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 diff --git a/atom/kv_transfer/offload/connector.py b/atom/kv_transfer/offload/connector.py index 36e1b66749..9c3c4052cb 100644 --- a/atom/kv_transfer/offload/connector.py +++ b/atom/kv_transfer/offload/connector.py @@ -335,6 +335,19 @@ 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_reports(self): # Contract is a 2-tuple (indexed, failed) -- matched by the kimi_k3 impl # and by the sole caller `indexed, failed = take()` in diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py index 89b254ba3f..86709051f5 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py +++ b/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 ( @@ -36,7 +37,10 @@ _JointPark, ) from atom.kv_transfer.offload.metadata import LMCacheOffloadMetadata -from atom.model_engine.state_offload import state_tier_cpu_bytes +from atom.model_engine.state_offload import ( + state_tier_budget_warning, + state_tier_cpu_bytes, +) from atom.utils import envs logger = logging.getLogger("atom") @@ -61,6 +65,11 @@ #: rather than leaving a partially-reported key pending forever. STATE_LOAD_DISPOSITION_CHANNEL = "k3_state_load_disposition" +#: One verdict per `StateTouchId`: the worker refreshed the image in LMCache's +#: LRU (`succeeded=True`) or found it gone (`succeeded=False`). Failure-dominant +#: across TP, like a joint load: one missing shard means the load would miss. +STATE_TOUCH_CHANNEL = "k3_state_touch" + #: 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 +124,8 @@ 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`. + self._touch_verdicts: dict[StateTouchId, bool] = {} def close(self) -> None: """Drain then join the state tier before the base executors. @@ -230,6 +241,7 @@ def _build_state_tier(self, transfer_tensors) -> None: staging_buffer_bytes=entry_bytes, release_after_transfer=gpu_connector.release_gpu_staging_after_transfer, ) + state_bytes = state_tier_cpu_bytes() codec = StateByteCodec( backend, staged, @@ -238,11 +250,17 @@ def _build_state_tier(self, transfer_tensors) -> None: world_size=world, worker_id=rank, layout_id=layout_id, - max_cache_bytes=state_tier_cpu_bytes(), + max_cache_bytes=state_bytes, ) # Share the allocator, but cap state retention separately so old # checkpoints cannot consume nearly the whole paged-KV CPU pool. codec.bind_storage_manager(self._engine.storage_manager) + pool_gib = float(getattr(cfg, "max_local_cpu_size", 0) or 0) + budget_warning = state_tier_budget_warning( + state_bytes, int(pool_gib * (1 << 30)), entry_bytes + ) + if budget_warning: + logger.warning("kimi_k3 offload: %s", budget_warning) # No index here: StateOffloadIndex lives in the engine process; both # directions report and the engine applies. self._state_tier = StateOffloadTier(codec) @@ -278,12 +296,42 @@ 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 state images the engine voted a hit for, and record + whether each is still there. + + Inline on this thread, not on the tier's executors: a touch is one + dict move under LocalCPU's lock, and it must not queue behind an + image-sized store. With no tier every image is gone by definition, so + the engine's index drops what it still advertises. + """ + 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 + continue + try: + resident = bool(tier.codec.touch(int(touch_id.prefix_hash))) + except Exception: + # Says nothing about the bytes; do not let it erase the index. + logger.warning( + "kimi_k3 offload: touch of state hash %d failed", + touch_id.prefix_hash, + exc_info=True, + ) + resident = True + self._touch_verdicts[touch_id] = resident + def _arm_joint_loads(self, metadata) -> None: """Hold a request owning both legs until both report. @@ -463,6 +511,13 @@ 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 in self._touch_verdicts.items(): + out.connector_completions.add( + ConnectorCompletion(STATE_TOUCH_CHANNEL, touch_id, resident) + ) + 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 +660,12 @@ 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] = [] + # Hit refresh: touches waiting for the next metadata, the round + # counter that keeps each one a distinct aggregator identity, and the + # verdicts drained by the engine (`take_state_touch_verdicts`). + self._pending_state_touches: list[StateTouchId] = [] + self._touch_round = 0 + self._state_touch_verdicts: dict[int, bool] = {} # 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 +703,25 @@ def enqueue_state_stores(self, stores) -> bool: self._pending_state_stores.extend(stores) return True + def enqueue_state_touches(self, hashes) -> bool: + """Queue LRU touches for state images the engine just voted a hit for. + + Each gets its own round, so a hash touched again is a new aggregator + identity rather than a tombstoned duplicate. + """ + if not hashes: + return False + for h in hashes: + self._pending_state_touches.append(StateTouchId(int(h), self._touch_round)) + self._touch_round += 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 build_connector_meta(self) -> LMCacheOffloadMetadata: self._refresh_save_stall() meta = super().build_connector_meta() @@ -654,6 +734,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 +761,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 -------------------------------------------------------- @@ -843,7 +926,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 @@ -855,6 +938,14 @@ 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_LOAD_DISPOSITION_CHANNEL: # Finding #6: a state load whose bytes survived a failed pair. Record # only the survivors -- their `failed_loading` report must abandon diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py index c9e19905b8..122859c7d9 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py @@ -65,13 +65,10 @@ def __init__( ) if self._max_entries is not None and self._max_entries < 1: raise ValueError("state CPU budget must fit at least one checkpoint") + # The cap's LRU, in the same order as LMCache's: refreshed by `get`, + # `touch` and a re-store, so the cap retires images nobody resumes + # from rather than the oldest stores. self._cache_keys: OrderedDict[Any, None] = OrderedDict() - # Under a budget, one extra reference per stored image. LMCache's - # LocalCPU LRU only evicts objects whose ref count is 1, so the pool - # cannot drop an image behind the codec's back while the engine-side - # `StateOffloadIndex` still advertises it. The budget is a reservation: - # only `put`'s own LRU retires these, KV pressure cannot. - self._held: dict[Any, Any] = {} self._cache_lock = threading.Lock() self._store_lock = threading.Lock() # Never hard-code a size: V4 keeps six compressor fields across @@ -122,37 +119,29 @@ def key(self, h: int): ) def put(self, h: int, unit_ids, on_source_released=None) -> bool: - # Bound state independently of KV while retaining the shared allocator. + # Cap state independently of KV while retaining the shared allocator; + # within the cap, images and KV chunks compete in LMCache's one LRU. # Only stores serialize; a load holds its MemoryObj reference while # unpacking, so removing its cache entry cannot free live bytes. if self._storage is None: return False + key = self.key(h) + if self._touch_key(key): + # Same prefix hash and layout, same bytes, still resident: storing + # again would only cost an image-sized D2H that LocalCPU then + # refuses as a duplicate key. The touch counts as the use it is. + if on_source_released is not None: + on_source_released() + return True if self._max_entries is None: return self._put(h, unit_ids, on_source_released) - key = self.key(h) with self._store_lock: with self._cache_lock: - if key in self._held: - # Same prefix hash and layout, same bytes, and the held - # reference keeps the stored copy resident: storing again - # would only cost an image-sized D2H that LocalCPU then - # refuses as a duplicate key. - self._cache_keys.move_to_end(key) - stored = True - else: - stored = False - if key not in self._cache_keys: - while len(self._cache_keys) >= self._max_entries: - old = next(iter(self._cache_keys)) - self._storage.remove(old) - del self._cache_keys[old] - held = self._held.pop(old, None) - if held is not None: - held.ref_count_down() - if stored: - if on_source_released is not None: - on_source_released() - return True + if key not in self._cache_keys: + while len(self._cache_keys) >= self._max_entries: + old = next(iter(self._cache_keys)) + self._storage.remove(old) + del self._cache_keys[old] ok = self._put(h, unit_ids, on_source_released) if ok: with self._cache_lock: @@ -160,6 +149,50 @@ def put(self, h: int, unit_ids, on_source_released=None) -> bool: self._cache_keys[key] = None return ok + def touch(self, h: int) -> bool: + """Refresh `h`'s image in LMCache's LRU; False when it is gone. + + The state-image twin of what a KV lookup does to its chunks + (`LocalCPUBackend.batched_async_contains(pin=True)` then + `touch_cache()` -> `cache_policy.update_on_hit`). State hits are + answered engine-side by `StateOffloadIndex`, so without this LMCache + never learns an image was used, and a prefix hit on every turn keeps + its KV at the young end and its image at the cold end, where pressure + evicts it first. + + `update_on_hit` under the backend's `cpu_lock`, not `contains(pin=True)` + + `touch_cache()`: a pin needs a matching unpin, and `keys_in_request` + is the list the KV lookup is filling at the same time. + """ + if self._storage is None: + return False + return self._touch_key(self.key(h)) + + def _touch_key(self, key) -> bool: + resident = self._refresh(key) + with self._cache_lock: + if resident: + if key in self._cache_keys: + self._cache_keys.move_to_end(key) + else: + # Evicted by LMCache: it no longer counts against the cap. + self._cache_keys.pop(key, None) + return resident + + def _refresh(self, key) -> bool: + backend = getattr(self._storage, "local_cpu_backend", None) + lock = getattr(backend, "cpu_lock", None) + hot = getattr(backend, "hot_cache", None) + policy = getattr(backend, "cache_policy", None) + if lock is not None and hot is not None and policy is not None: + with lock: + if key in hot: + policy.update_on_hit(key, hot) + return True + # Not in the hot cache (or there is none): another backend may still + # hold it, and `get` reads every backend. Presence only. + return bool(self._storage.contains(key)) + def _put(self, h: int, unit_ids, on_source_released=None) -> bool: """Store one checkpoint image. False when nothing was stored. @@ -186,9 +219,6 @@ def _put(self, h: int, unit_ids, on_source_released=None) -> bool: obj = self._allocate(self.entry_bytes) if obj is None: return False - hold = self._max_entries is not None - if hold: - obj.ref_count_up() # `batched_put` discharges the reference it is handed, but only in its # terminal tail loop (LMCache `storage_manager.py`, one # `ref_count_down()` per object), which runs *after* every step that can @@ -221,19 +251,7 @@ def _put(self, h: int, unit_ids, on_source_released=None) -> bool: self._storage.batched_put([key], [obj]) except Exception: obj.ref_count_down() - if hold: - obj.ref_count_down() raise - if hold: - ref_count = getattr(obj, "get_ref_count", None) - if not callable(ref_count) or ref_count() > 1: - # The hot cache adopted it (its own reference plus ours). - with self._cache_lock: - self._held[key] = obj - else: - # Refused, e.g. the key was already resident: ours is the - # last reference, and keeping it would strand the block. - obj.ref_count_down() return True def get(self, h: int, slot) -> bool: @@ -251,8 +269,11 @@ def get(self, h: int, slot) -> bool: """ if self._storage is None: return False - obj = self._storage.get(self.key(h)) + key = self.key(h) + obj = self._storage.get(key) if obj is None: + with self._cache_lock: + self._cache_keys.pop(key, None) return False size = self._object_bytes(obj) if size is not None and size != self.entry_bytes: @@ -276,10 +297,8 @@ def get(self, h: int, slot) -> bool: self._staged.unpack(obj, self._backend.state_entry_views(slot)) finally: obj.ref_count_down() - with self._cache_lock: - key = self.key(h) - if key in self._cache_keys: - self._cache_keys.move_to_end(key) + # `get_blocking` does not refresh LMCache's LRU; a resume is a use. + self._touch_key(key) return True @staticmethod diff --git a/atom/kv_transfer/offload/metadata.py b/atom/kv_transfer/offload/metadata.py index 3fd06b6bcf..d54a40498c 100644 --- a/atom/kv_transfer/offload/metadata.py +++ b/atom/kv_transfer/offload/metadata.py @@ -197,6 +197,7 @@ class LMCacheOffloadMetadata(ConnectorMetadata): "lookup_requests_in_step", "state_loads", "state_stores", + "state_touches", ) def __init__(self) -> None: @@ -215,6 +216,11 @@ def __init__(self) -> None: # generation) is what the source-release and index reports carry back to # settle the pin and index the hash. self.state_stores: list[tuple] = [] + # StateTouchId per state image the engine voted a hit for: each worker + # refreshes the image in LMCache's LRU and reports whether it is still + # there. Work in its own right -- an admission can be waiting on the + # verdict. + self.state_touches: list = [] def add_request(self, meta: LMCacheReqMeta) -> None: self.requests.append(meta) diff --git a/atom/model_engine/block_manager.py b/atom/model_engine/block_manager.py index e5cbce1004..f5f7d29b83 100644 --- a/atom/model_engine/block_manager.py +++ b/atom/model_engine/block_manager.py @@ -498,6 +498,9 @@ def __init__( if checkpoint_spec is not None else None ), + # Off again (`stop_refreshing`) if the connector turns out not + # to carry state touches. + refresh_hits=True, ) # Attached rather than passed at construction: the coordinator is # built before the switch is read (it needs `checkpoint_spec`, @@ -1434,12 +1437,68 @@ def _walk_still_valid(self, seq: Sequence, memo: _CachedWalk) -> bool: ) def record_allocation( - self, seq: Sequence, num_cached_blocks: int, block_hashes: list[int] - ) -> None: - """Commit the fit probe before any allocation changes the pool.""" - if self.enable_prefix_caching: - decision = self._joint_kv_boundary(seq, num_cached_blocks, block_hashes) - self._commit_joint_boundary(seq, decision) + self, + seq: Sequence, + num_cached_blocks: int, + block_hashes: list[int], + *, + defer_unverified: bool = False, + ) -> bool: + """Commit the fit probe before any allocation changes the pool. + + With `defer_unverified`, returns False and commits nothing when the + state this admission would load from the CPU tier has not been + confirmed in LMCache lately (`StateOffloadIndex.needs_verify`). The + caller holds the request for a pass while that image's touch comes + back, rather than parking a joint load that may miss and recompute the + whole prompt. + """ + if not self.enable_prefix_caching: + return True + decision = self._joint_kv_boundary(seq, num_cached_blocks, block_hashes) + if defer_unverified and self._tier_state_unverified( + seq, num_cached_blocks, block_hashes, decision + ): + self.state_offload.loads_deferred_for_verify += 1 + return False + self._commit_joint_boundary(seq, decision) + return True + + def _tier_state_unverified( + self, + seq: Sequence, + num_cached_blocks: int, + block_hashes: list[int], + decision: JointBoundaryDecision, + ) -> bool: + """Whether `allocate` would load this seq's state from the CPU tier + under a hash not confirmed resident recently. + + Mirrors which hash `allocate` hands `_attach_state_slots`: the joint + boundary when there is one, else the HBM hit's last block; and the + tier only serves it when HBM does not hold it. + """ + offload = self.state_offload + if offload is None or not offload.refresh_hits or not seq.has_per_req_cache: + return False + if decision.skip_reason is None: + if decision.state_from_hbm: + return False + h = decision.boundary_hash + else: + if not 0 < num_cached_blocks <= len(block_hashes): + return False + if self.state.transfer.replays or ( + self.state_replay_tokens and seq.replays_hit + ): + return False # a replay rebuilds the state; nothing is loaded + h = block_hashes[num_cached_blocks - 1] + if self.paged_state_checkpoints is not None: + if self.paged_state_checkpoints.contains(h): + return False + elif self.state.lookup(h) >= 0: + return False + return offload.needs_verify(h) def allocate(self, seq: Sequence, num_cached_blocks: int = 0) -> bool: """Allocate blocks for `seq`. `num_cached_blocks` is the hit count diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index 825ecabb1c..d134d373a5 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -2288,7 +2288,15 @@ def _schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]] | None: ): self.waiting.appendleft(seq) break - self.block_manager.record_allocation(seq, num_cached_blocks, block_hashes) + if not self.block_manager.record_allocation( + seq, num_cached_blocks, block_hashes, defer_unverified=True + ): + # The CPU state image this admission would load has not been + # confirmed in LMCache lately and its touch is out. Held like + # an in-flight prefix: an image found gone is forgotten when + # the verdict lands, and the retry plans without it. + prefix_waiters.append(seq) + continue if not self.block_manager.allocate(seq, num_cached_blocks): # A state-less joint boundary could not be privatised without # risking a shared decoding sequence's blocks (finding #2). @@ -3649,6 +3657,39 @@ def _publish_state_stores(self) -> None: # never reached a worker, so it is not a `stores_failed`. self.block_manager.settle_state_store(op, ok=False, attempted=False) + def _publish_state_touches(self) -> None: + """Hand the state images this pass voted a hit for to the connector, + which refreshes them in LMCache's LRU and reports whether each is + still there (`_apply_state_touch_verdicts`). + + A connector that will not carry them turns the refresh off for good: + otherwise every admission that waits for a verdict would wait out + `OFFLOAD_LOOKUP_DEFER_S` for nothing. + """ + offload = getattr(getattr(self, "block_manager", None), "state_offload", None) + if offload is None or not offload.refresh_hits: + return + hashes = offload.take_touches() + if not hashes: + return + enqueue = getattr(self.kv_connector, "enqueue_state_touches", None) + if enqueue is not None and enqueue(hashes): + return + logger.warning( + "state offload: %s does not carry state touches; state images " + "will not be refreshed in LMCache's LRU, and an image LMCache " + "evicts is only noticed when a load misses.", + type(self.kv_connector).__name__, + ) + offload.stop_refreshing() + + def _apply_state_touch_verdicts(self, offload) -> None: + take = getattr(self.kv_connector, "take_state_touch_verdicts", None) + if take is None: + return + for h, resident in take().items(): + offload.apply_touch_verdict(h, resident) + def _publish_state_loads(self) -> None: """Hand this pass's state-tier loads to the connector, before it builds its metadata. @@ -3659,6 +3700,7 @@ def _publish_state_loads(self) -> None: """ if self.kv_connector is None: return + self._publish_state_touches() # Guarded like the `_settle_state_load` / `_publish_state_stores` # siblings: a scheduler built without a block manager (the connector # test doubles) would otherwise AttributeError on the bare deref. @@ -4776,6 +4818,9 @@ def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput): # against a `get` that must miss. bm = getattr(self, "block_manager", None) offload = getattr(bm, "state_offload", None) + if offload is not None: + # Before the next admission pass, which reads what they decide. + self._apply_state_touch_verdicts(offload) take = getattr(self.kv_connector, "take_state_reports", None) if offload is not None and take is not None: indexed, failed = take() diff --git a/atom/model_engine/state_offload.py b/atom/model_engine/state_offload.py index 4ed4f1bc1c..9a4742d2bc 100644 --- a/atom/model_engine/state_offload.py +++ b/atom/model_engine/state_offload.py @@ -15,18 +15,36 @@ import os from collections import OrderedDict from dataclasses import dataclass, replace +from time import monotonic logger = logging.getLogger("atom") +#: How long a hash counts as known-resident after the last evidence that its +#: image was in LMCache (stored, loaded, or confirmed by a touch). Every one of +#: those leaves the image at the young end of LMCache's LRU, and the pool turns +#: over in minutes, so this is far inside the time an image needs to age out. +_TIER_VERIFIED_FOR_S = 5.0 + +#: A hash voted on every scheduling pass is touched at most this often. The +#: touch only has to keep the image near the young end, not track every pass. +_TIER_TOUCH_INTERVAL_S = 0.5 + class StateOffloadIndex: """What is believed to be in LMCache, and what is being fetched back. - `hashes` answers the membership half of the tier's vote. It is deliberately - optimistic: LMCache's own LRU can drop bytes at any time, so a hash here - means "was stored once", never "is still there". The false positive costs - one lookup and a park/unpark, and is handled by the `failed_loading` path - (`fail_load` -> `forget`). + `hashes` answers the membership half of the tier's vote. LMCache's own LRU + can drop bytes at any time, so on its own a hash here means "was stored + once", never "is still there". + + With `refresh_hits`, every hit the index votes for is also *touched*: the + hash is queued for `take_touches`, the worker refreshes the image's + position in LMCache's LRU the same way a KV lookup refreshes its chunks, + and reports whether the image was still there (`apply_touch_verdict`). That + keeps a prefix's state image as warm as its KV, and lets a hash LMCache + dropped leave the index before a request parks on it (`needs_verify`). + A false positive that slips through still costs one park/unpark and is + handled by the `failed_loading` path (`fail_load` -> `forget`). """ def __init__( @@ -35,6 +53,7 @@ def __init__( can_store: bool = True, can_load: bool = True, max_cpu_entries: int | None = None, + refresh_hits: bool = False, ) -> None: # The two legs are separately granted. A `kv_producer` role saves and # never loads, a `kv_consumer` loads and never saves, and offering a @@ -86,6 +105,23 @@ def __init__( # than indexed. Any non-zero value means the reclaim window is firing # on live transfers and is set too low. self.stores_untrusted = 0 + # Hit refresh (see the class docstring). Off for indexes whose + # connector carries no touches (the vLLM plugin), where the queue would + # only grow. + self.refresh_hits = bool(refresh_hits) + # Hashes voted for since the last `take_touches`. + self._touch_wanted: set[int] = set() + # hash -> monotonic of the last touch handed out / the last evidence + # the image was resident. + self._touch_sent_at: dict[int, float] = {} + self._verified_at: dict[int, float] = {} + self.touches_sent = 0 + # Touches that found the image gone: LMCache evicted it and the index + # learned of it before any request parked on it. + self.touch_misses = 0 + # Admissions held back a pass because the image they would load had + # not been confirmed recently (`needs_verify`). + self.loads_deferred_for_verify = 0 # ------------------------------- stores -------------------------------- # def note_stored(self, h: int) -> None: @@ -101,20 +137,34 @@ def note_stored(self, h: int) -> None: # re-stored prefix is treated as freshly used, then trim the cold end. self._hash_lru.pop(h, None) self._hash_lru[h] = None + self._verified_at[h] = monotonic() while len(self._hash_lru) > self._hash_cap: old, _ = self._hash_lru.popitem(last=False) self.hashes.discard(old) + self._drop_touch_state(old) self.hashes_evicted += 1 def forget(self, h: int) -> None: - """Drop a hash whose load failed, so the next request does not retry.""" + """Drop a hash whose image is gone, so the next request does not retry.""" self.hashes.discard(h) self._hash_lru.pop(h, None) + self._drop_touch_state(h) + + def _drop_touch_state(self, h: int) -> None: + self._touch_wanted.discard(h) + self._touch_sent_at.pop(h, None) + self._verified_at.pop(h, None) def could_serve(self, h: int) -> bool: """Whether a load for `h` could be offered at all: this tier can load AND believes it holds `h`. + A True answer is a hit vote, so under `refresh_hits` it also queues `h` + for a touch. This is the one predicate every voter goes through, which + is what makes it the place to refresh: a prefix whose state is voted + for on every pass stays as warm in LMCache as its KV chunks, which the + KV lookup refreshes on every pass. + The membership+capability half of `request_load`'s guard (minus its per-request in-flight check), factored out so the admission-path voters -- `PageUnitCheckpointCoordinator._reachable` and @@ -126,7 +176,85 @@ def could_serve(self, h: int) -> bool: still-resident HBM rung, and `request_load` returned False -- the boundary disowned and the HBM checkpoint forfeited to a full recompute. """ - return self.can_load and h in self.hashes + if not (self.can_load and h in self.hashes): + return False + if self.refresh_hits: + self._touch_wanted.add(h) + return True + + # ------------------------------- touches ------------------------------- # + def take_touches(self) -> list[int]: + """Hashes to touch in LMCache this step, drained. + + A hash already touched within `_TIER_TOUCH_INTERVAL_S` is skipped: a + request waiting for admission votes for the same hash on every pass. + """ + if not self._touch_wanted: + return [] + now = monotonic() + out = [] + for h in self._touch_wanted: + if now - self._touch_sent_at.get(h, -1e18) < _TIER_TOUCH_INTERVAL_S: + continue + self._touch_sent_at[h] = now + out.append(h) + self._touch_wanted.clear() + self.touches_sent += len(out) + return out + + def apply_touch_verdict(self, h: int, resident: bool) -> None: + """A touch came back: the image was still in LMCache, or it was not. + + Gone means LMCache's LRU dropped it, so the hash leaves the index now, + before a request parks a joint load on it and recomputes the whole + prompt after the miss. + """ + h = int(h) + if h not in self.hashes: + return + if not resident: + self.touch_misses += 1 + self.forget(h) + return + self._verified_at[h] = monotonic() + # The touch made it LMCache's youngest image; keep the cap's order. + self._hash_lru.move_to_end(h) + + def stop_refreshing(self) -> None: + """The connector cannot carry touches: stop queueing them and stop + holding admissions for verdicts that will never come.""" + self.refresh_hits = False + self._touch_wanted.clear() + self._touch_sent_at.clear() + + def needs_verify(self, h: int) -> bool: + """Whether an admission about to load `h` should wait a pass for a + touch verdict instead of parking on it now. + + True while `h` has not been confirmed resident within + `_TIER_VERIFIED_FOR_S` and its touch is still pending. A pass later + the verdict is normally in: resident lets the load go ahead, gone has + already dropped `h`, so the admission picks a lower boundary (an HBM + checkpoint, an older CPU image, or a plain prefill) instead of parking + and recomputing everything after the miss. A verdict that has not come + back within `defer_s` of its touch stops holding the admission, which + then proceeds as it did before touches existed. + """ + if not self.refresh_hits or h not in self.hashes: + return False + defer_s = verify_defer_s() + if defer_s <= 0: + return False + now = monotonic() + verified = self._verified_at.get(h) + if verified is not None and now - verified <= _TIER_VERIFIED_FOR_S: + return False + sent = self._touch_sent_at.get(h) + if sent is None or (verified is not None and verified >= sent): + # No touch outstanding: ask for one; it leaves this pass. + self._touch_wanted.add(h) + return True + return now - sent <= defer_s # -------------------------------- loads -------------------------------- # def request_load(self, req_id, h: int) -> bool: @@ -169,6 +297,7 @@ def complete_load(self, req_id) -> None: # order so the cap trims what the codec will trim. if h in self._hash_lru: self._hash_lru.move_to_end(h) + self._verified_at[h] = monotonic() def fail_load(self, req_id) -> None: """No bytes came back. Retract the claim as well as counting it. @@ -220,9 +349,28 @@ def stats(self) -> dict[str, int]: # far past `indexed` means the state budget is smaller than the # reusable working set. "hashes_evicted": self.hashes_evicted, + # Hit refresh. `touch_misses` are images LMCache dropped that left + # the index before anything parked on them; `loads_failed` is what + # still got past. + "touches_sent": self.touches_sent, + "touch_misses": self.touch_misses, + "loads_deferred_for_verify": self.loads_deferred_for_verify, } +def verify_defer_s() -> float: + """How long an admission may wait for a touch verdict + (`StateOffloadIndex.needs_verify`). + + `OFFLOAD_LOOKUP_DEFER_S`: the same question it already answers for a KV + lookup in flight -- how long admission may pass over a request whose tier + answer is on its way rather than block on it or guess. + """ + from atom.utils import envs + + return float(envs.OFFLOAD_LOOKUP_DEFER_S) + + #: Offload layouts whose worker half builds a `StateOffloadTier`. Not every #: `lmcache_offload` does: `dense` has no per-request state at all and `hybrid` #: (DSV4) keeps its own in the SLOT sidecar. `qwen` subclasses the K3 tier. @@ -388,7 +536,8 @@ def state_tier_capability(config) -> StateTierCapability: def state_tier_cpu_bytes() -> int: """CPU bytes the state tier may keep per worker (`OFFLOAD_STATE_CPU_SIZE`, - GiB, default 32), carved from the paged-KV LMCache pool it shares. + GiB, default 32): a cap on the images in the paged-KV LMCache pool they + share, not a reservation -- images and KV chunks compete in one LRU. One reader for both halves: the worker's `StateByteCodec` evicts against it and the engine's `StateOffloadIndex` mirrors the resulting entry bound. @@ -396,6 +545,41 @@ def state_tier_cpu_bytes() -> int: return int(float(os.environ.get("OFFLOAD_STATE_CPU_SIZE", "32")) * (1 << 30)) +def state_tier_budget_warning( + state_bytes: int, pool_bytes: int, entry_bytes: int +) -> str | None: + """Why the state cap and the LMCache pool size do not fit together, or None. + + The cap only bounds how much of the shared pool images may take, so a cap + at or above the pool bounds nothing, and one above half of it lets images + push most of the KV out under pressure. Neither stops a store -- the pool + evicts by recency either way -- but both are almost always a sizing + mistake (e.g. `LMCACHE_MAX_LOCAL_CPU_SIZE=32` with the 32 GiB default cap), + so the caller says so loudly at startup. + """ + if pool_bytes <= 0 or state_bytes <= 0: + return None + gib = float(1 << 30) + sizes = ( + f"OFFLOAD_STATE_CPU_SIZE={state_bytes / gib:.1f} GiB " + f"({state_bytes // max(1, entry_bytes)} images of " + f"{entry_bytes / (1 << 20):.1f} MiB), " + f"LMCACHE_MAX_LOCAL_CPU_SIZE={pool_bytes / gib:.1f} GiB" + ) + if state_bytes >= pool_bytes: + return ( + f"{sizes}: the state cap is not below the pool, so it bounds " + "nothing and state images can take the whole pool from paged KV. " + "Lower OFFLOAD_STATE_CPU_SIZE or raise LMCACHE_MAX_LOCAL_CPU_SIZE." + ) + if 2 * state_bytes > pool_bytes: + return ( + f"{sizes}: state images may take more than half of the pool, " + "leaving paged KV the smaller share under pressure." + ) + return None + + def state_tier_chunk_tokens(config) -> int: """The LMCache chunk size in tokens for the joint KV leg, or 0. diff --git a/docs/environment_variables.md b/docs/environment_variables.md index e554b98d1b..456e9cb3a2 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -386,7 +386,7 @@ derive a bound. | **OFFLOAD_LOOKUP_MEMO_STEPS** | int | 32 | Scheduler steps a memoised tier-lookup answer is replayed before the tier is asked again. | | **OFFLOAD_LOOKUP_RETRY_STEPS** | int | 32 | Scheduler steps a failed tier lookup suppresses the next attempt. | | **OFFLOAD_ASYNC_LOOKUP_DEPTH** | int | 16 | `lmcache_mp` sends tier lookups for the head of the waiting queue before admission, as non-blocking requests, and admission consumes the answer; a request whose answer has not arrived is passed over (in arrival order) instead of blocking the scheduler. This is how many waiting requests per scheduling pass get a prefetched lookup. | -| **OFFLOAD_LOOKUP_DEFER_S** | float | 2.0 | How long `lmcache_mp` admission may pass over a request whose prefetched lookup is still in flight before it waits for the answer (finite, ≥ 0). | +| **OFFLOAD_LOOKUP_DEFER_S** | float | 2.0 | How long `lmcache_mp` admission may pass over a request whose prefetched lookup is still in flight before it waits for the answer (finite, ≥ 0). Also bounds how long `lmcache_offload` with a state tier (Kimi-K3, Qwen) holds an admission that would load a CPU state image not confirmed recently, while that image's LRU touch reports whether LMCache still holds it; `0` disables the hold. | | **OFFLOAD_SKIP_LOOKUP_ON_HBM_HIT** | bool | 1 (true) | On by default; set to `0` to always wait for the lookup. Admit a prefill without waiting for its external-tier lookup when the lookup cannot change what is loaded: the HBM prefix leaves fewer than `OFFLOAD_MIN_LOAD_TOKENS` tokens uncached, so every load path would refuse the tier hit. Requires an asynchronous lookup (LMCache MP) and is not applied with joint KV/state offload. The lookup is still sent, because it refreshes the prefix's L1 LRU recency and opens the MP session whose end touches the request's chunks; its answer is not used as a load hit. The save floor starts at the HBM prefix and is lowered to the tier's contiguous hit when the answer arrives (or to 0 if it never does), so chunks the tier evicted are stored again. Read locks are released when the answer arrives; the MP session ends after that, and an unanswered lookup is given up after 120 s, leaving its locks to the read TTL. | | **OFFLOAD_PROFILE** | bool | 0 | Emit `[OFFLOAD-SAVE-PROF]` / `[OFFLOAD-LOAD-PROF]` per-transfer records, and once a minute an `[OFFLOAD-PERF]` line per scheduler with lookup/save/load-skip counters and per-step timings. An empty value reads as off. | | **OFFLOAD_SINGLE_STREAM** | bool | 0 | Experimental: run the staging pack and copy legs on one stream. | diff --git a/tests/test_lmcache_offload_connector.py b/tests/test_lmcache_offload_connector.py index c4e80b7d54..9c68522edc 100644 --- a/tests/test_lmcache_offload_connector.py +++ b/tests/test_lmcache_offload_connector.py @@ -5822,6 +5822,9 @@ def _k3_scheduler() -> KimiK3OffloadScheduler: s._save_rr_last = None s._pending_state_loads = [] s._pending_state_stores = [] + s._pending_state_touches = [] + s._touch_round = 0 + s._state_touch_verdicts = {} s._save_inflight_since = {} s._save_stalled = False s._warned_save_stalled = False diff --git a/tests/test_state_tier_lru_refresh.py b/tests/test_state_tier_lru_refresh.py new file mode 100644 index 0000000000..709da83582 --- /dev/null +++ b/tests/test_state_tier_lru_refresh.py @@ -0,0 +1,687 @@ +# SPDX-License-Identifier: MIT +"""State images are refreshed in LMCache's LRU like the KV they belong to. + +State images and paged KV share one LocalCPU pool. A KV lookup refreshes every +chunk it touches (`contains(pin=True)` then `touch_cache()` -> +`cache_policy.update_on_hit`), but a state hit is answered engine-side by +`StateOffloadIndex`, so LMCache never learned an image was used: a prefix hit +on every turn kept its KV at the young end and its image at the cold end, where +pool pressure evicted it first, and the index kept advertising it until a joint +load parked on it, missed, and recomputed the whole prompt. + +These tests model LocalCPU's hot cache and LRU policy, and cover the three +pieces of the fix: the codec touch, the engine-side touch/verdict loop that +keeps the index in step with evictions, and the admission hold that lets a +verdict land before a request parks on an image. +""" + +import dataclasses +import sys +import threading +import types +from collections import OrderedDict +from types import SimpleNamespace + +import pytest + +from atom.kv_transfer.disaggregation.types import ConnectorCompletion, StateTouchId +from atom.model_engine import state_offload as so +from atom.model_engine.state_offload import ( + StateOffloadIndex, + state_tier_budget_warning, +) + + +@pytest.fixture +def lmcache_key(monkeypatch): + @dataclasses.dataclass(frozen=True) + class CacheEngineKey: + model_name: str + world_size: int + worker_id: int + chunk_hash: int + dtype: object + + utils_module = types.ModuleType("lmcache.utils") + utils_module.CacheEngineKey = CacheEngineKey + lmcache_module = types.ModuleType("lmcache") + lmcache_module.__path__ = [] + lmcache_module.utils = utils_module + monkeypatch.setitem(sys.modules, "lmcache", lmcache_module) + monkeypatch.setitem(sys.modules, "lmcache.utils", utils_module) + return CacheEngineKey + + +@pytest.fixture +def clock(monkeypatch): + """`state_offload`'s monotonic clock, advanced by hand.""" + now = [1000.0] + monkeypatch.setattr(so, "monotonic", lambda: now[0]) + return now + + +class _Obj: + """A MemoryObj's reference count, and nothing else.""" + + def __init__(self, size: int = 1024) -> None: + self.refs = 1 + self.size = size + + def ref_count_up(self) -> None: + self.refs += 1 + + def ref_count_down(self) -> None: + assert self.refs > 0, "reference discharged twice" + self.refs -= 1 + + def get_ref_count(self) -> int: + return self.refs + + def get_size(self) -> int: + return self.size + + +class _LRUPolicy: + """LMCache's `LRUCachePolicy.update_on_hit`.""" + + def update_on_hit(self, key, cache) -> None: + cache.move_to_end(key) + + +class _LocalCPU: + """`storage_manager` over LocalCPU's rules. + + `batched_put` adopts an object with its own reference unless the key is + already resident (LocalCPU refuses duplicates), then drops the caller's. + `get` does not touch the LRU (`get_blocking` only takes a reference). + `evict_oldest` is the pool making room for KV: oldest first, skipping + anything another holder still references. + """ + + def __init__(self, refuse_all: bool = False) -> None: + self.hot: OrderedDict = OrderedDict() + self.refuse_all = refuse_all + self.fail_put = False + self.local_cpu_backend = SimpleNamespace( + cpu_lock=threading.Lock(), hot_cache=self.hot, cache_policy=_LRUPolicy() + ) + + def batched_put(self, keys, objs) -> None: + if self.fail_put: + raise RuntimeError("put failed") + for k, o in zip(keys, objs, strict=True): + if not self.refuse_all and k not in self.hot: + o.ref_count_up() + self.hot[k] = o + o.ref_count_down() + + def get(self, key): + o = self.hot.get(key) + if o is not None: + o.ref_count_up() + return o + + def contains(self, key): + return "LocalCPUBackend" if key in self.hot else None + + def remove(self, key) -> int: + o = self.hot.pop(key, None) + if o is None: + return 0 + o.ref_count_down() + return 1 + + def evict_oldest(self): + for k, o in self.hot.items(): + if o.get_ref_count() == 1: + self.hot.pop(k).ref_count_down() + return k + return None + + +def _codec(storage, max_entries): + from atom.kv_transfer.offload.hybrid.kimi_k3.state_object import StateByteCodec + + codec = StateByteCodec( + backend=SimpleNamespace( + page_unit_views=lambda units: [], state_entry_views=lambda slot: [] + ), + staged=SimpleNamespace(pack=lambda views, obj: None, unpack=lambda o, v: None), + entry_bytes=1024, + model_name="qwen", + world_size=1, + worker_id=0, + layout_id="L", + max_cache_bytes=None if max_entries is None else max_entries * 1024, + ) + codec.bind_storage_manager(storage) + allocated = [] + + def allocate(nbytes): + allocated.append(_Obj(nbytes)) + return allocated[-1] + + codec._allocate = allocate + return codec, allocated + + +def _kv(pool, name): + pool.hot[name] = _Obj() + + +# ── the codec: touch, get and re-store refresh LMCache's LRU ─────────────── + + +class TestCodecRefresh: + def test_an_untouched_image_is_the_first_thing_pressure_evicts(self, lmcache_key): + """The bug: KV written after the image is younger, however often the + prefix hits.""" + pool = _LocalCPU() + codec, _ = _codec(pool, max_entries=None) + assert codec.put(1, []) + _kv(pool, "kv-chunk") + assert pool.evict_oldest() == codec.key(1) + + def test_a_touch_moves_the_image_past_older_kv(self, lmcache_key): + pool = _LocalCPU() + codec, _ = _codec(pool, max_entries=None) + assert codec.put(1, []) + _kv(pool, "kv-chunk") + assert codec.touch(1) is True + assert pool.evict_oldest() == "kv-chunk" + assert codec.key(1) in pool.hot + + def test_a_touch_reports_an_evicted_image_gone(self, lmcache_key): + pool = _LocalCPU() + codec, _ = _codec(pool, max_entries=2) + assert codec.put(1, []) + assert pool.evict_oldest() == codec.key(1) + assert codec.touch(1) is False + assert codec.key(1) not in codec._cache_keys, "no longer counts against the cap" + + def test_a_resume_refreshes_the_image(self, lmcache_key): + """`get_blocking` takes a reference and nothing else.""" + pool = _LocalCPU() + codec, objs = _codec(pool, max_entries=None) + assert codec.put(1, []) + _kv(pool, "kv-chunk") + assert codec.get(1, 0) + assert objs[0].refs == 1, "the load's reference is discharged" + assert pool.evict_oldest() == "kv-chunk" + + def test_a_missed_get_leaves_the_cap_lru(self, lmcache_key): + pool = _LocalCPU() + codec, _ = _codec(pool, max_entries=2) + assert codec.put(1, []) + pool.evict_oldest() + assert not codec.get(1, 0) + assert codec._cache_keys == OrderedDict() + + def test_a_resident_image_is_refreshed_not_stored_twice(self, lmcache_key): + pool = _LocalCPU() + codec, objs = _codec(pool, max_entries=2) + released = [] + assert codec.put(1, []) + _kv(pool, "kv-chunk") + assert codec.put(1, [], on_source_released=lambda: released.append(True)) + + assert len(objs) == 1, "no second image-sized allocation and D2H" + assert released == [True], "the source units must still be handed back" + assert pool.evict_oldest() == "kv-chunk" + + def test_an_evicted_image_is_stored_again(self, lmcache_key): + pool = _LocalCPU() + codec, objs = _codec(pool, max_entries=2) + assert codec.put(1, []) + pool.evict_oldest() + assert codec.put(1, []) + assert len(objs) == 2 and codec.key(1) in pool.hot + + def test_the_cap_retires_the_least_recently_used_image(self, lmcache_key): + """Not the oldest store: the image a touch kept warm survives.""" + pool = _LocalCPU() + codec, objs = _codec(pool, max_entries=2) + assert codec.put(1, []) and codec.put(2, []) + assert codec.touch(1) + assert codec.put(3, []) + + assert codec.key(2) not in pool.hot + assert objs[1].refs == 0, "a retired image goes back to the allocator" + assert list(codec._cache_keys) == [codec.key(1), codec.key(3)] + + def test_images_hold_no_extra_reference(self, lmcache_key): + """No reservation: KV pressure may evict an image like any chunk.""" + pool = _LocalCPU() + codec, objs = _codec(pool, max_entries=2) + assert codec.put(1, []) + assert objs[0].refs == 1 + assert pool.evict_oldest() == codec.key(1) + + def test_a_refused_put_keeps_no_reference(self, lmcache_key): + pool = _LocalCPU(refuse_all=True) + codec, objs = _codec(pool, max_entries=2) + assert codec.put(1, []) + assert objs[0].refs == 0, "holding a refused object would strand its block" + + def test_a_failed_put_discharges_its_reference(self, lmcache_key): + pool = _LocalCPU() + pool.fail_put = True + codec, objs = _codec(pool, max_entries=2) + with pytest.raises(RuntimeError): + codec.put(1, []) + assert objs[0].refs == 0 + + def test_without_a_local_cpu_backend_a_touch_asks_presence(self, lmcache_key): + pool = _LocalCPU() + codec, _ = _codec(pool, max_entries=None) + assert codec.put(1, []) + pool.local_cpu_backend = None + assert codec.touch(1) is True + pool.hot.clear() + assert codec.touch(1) is False + + +# ── the index: votes queue touches, verdicts keep it in step ─────────────── + + +def _index(*hashes, refresh=True): + idx = StateOffloadIndex(refresh_hits=refresh) + for h in hashes: + idx.note_stored(h) + return idx + + +class TestIndexTouches: + def test_a_hit_vote_queues_a_touch(self, clock): + idx = _index(7) + assert idx.could_serve(7) and not idx.could_serve(8) + assert idx.take_touches() == [7] + assert idx.take_touches() == [] + + def test_no_touches_without_refresh(self, clock): + idx = _index(7, refresh=False) + assert idx.could_serve(7) + assert idx.take_touches() == [] + + def test_a_hash_voted_every_pass_is_touched_at_an_interval(self, clock): + idx = _index(7) + idx.could_serve(7) + assert idx.take_touches() == [7] + clock[0] += so._TIER_TOUCH_INTERVAL_S / 2 + idx.could_serve(7) + assert idx.take_touches() == [] + clock[0] += so._TIER_TOUCH_INTERVAL_S + idx.could_serve(7) + assert idx.take_touches() == [7] + assert idx.touches_sent == 2 + + def test_a_gone_verdict_drops_the_hash(self, clock): + idx = _index(7) + idx.apply_touch_verdict(7, False) + assert not idx.could_serve(7) + assert idx.touch_misses == 1 + assert idx.stats()["touch_misses"] == 1 + + def test_a_resident_verdict_refreshes_the_cap_order(self, clock): + idx = StateOffloadIndex(max_cpu_entries=19, refresh_hits=True) # cap 3 + for h in (1, 2, 3): + idx.note_stored(h) + idx.apply_touch_verdict(1, True) + idx.note_stored(4) + assert idx.could_serve(1) and not idx.could_serve(2) + + def test_a_verdict_for_a_forgotten_hash_is_ignored(self, clock): + idx = _index(7) + idx.forget(7) + idx.apply_touch_verdict(7, True) + assert not idx.could_serve(7) + + def test_stop_refreshing_drops_pending_touches(self, clock): + idx = _index(7) + idx.could_serve(7) + idx.stop_refreshing() + assert idx.take_touches() == [] + assert not idx.needs_verify(7) + + +class TestNeedsVerify: + def test_a_fresh_store_needs_no_verdict(self, clock): + assert not _index(7).needs_verify(7) + + def test_a_stale_hash_waits_for_its_touch(self, clock): + idx = _index(7) + clock[0] += so._TIER_VERIFIED_FOR_S + 1 + assert idx.needs_verify(7) + assert idx.take_touches() == [7], "the hold asks for its own touch" + assert idx.needs_verify(7), "still out" + idx.apply_touch_verdict(7, True) + assert not idx.needs_verify(7) + + def test_a_gone_verdict_ends_the_hold_by_dropping_the_hash(self, clock): + idx = _index(7) + clock[0] += so._TIER_VERIFIED_FOR_S + 1 + assert idx.needs_verify(7) + idx.take_touches() + idx.apply_touch_verdict(7, False) + assert not idx.needs_verify(7) and not idx.could_serve(7) + + def test_a_lost_verdict_holds_no_longer_than_the_lookup_defer( + self, clock, monkeypatch + ): + monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "2") + idx = _index(7) + clock[0] += so._TIER_VERIFIED_FOR_S + 1 + assert idx.needs_verify(7) + idx.take_touches() + clock[0] += 1.5 + assert idx.needs_verify(7) + clock[0] += 1.0 + assert not idx.needs_verify(7) + + def test_a_zero_defer_never_holds(self, clock, monkeypatch): + monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "0") + idx = _index(7) + clock[0] += so._TIER_VERIFIED_FOR_S + 1 + assert not idx.needs_verify(7) + + def test_a_completed_load_counts_as_verified(self, clock): + idx = _index(7) + clock[0] += so._TIER_VERIFIED_FOR_S + 1 + assert idx.request_load("r", 7) + idx.complete_load("r") + assert not idx.needs_verify(7) + + +# ── the connector: touches out in the metadata, verdicts back by quorum ──── + + +def _k3_worker(tier): + from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( + KimiK3OffloadConnector, + ) + + c = KimiK3OffloadConnector.__new__(KimiK3OffloadConnector) + c._state_tier = tier + c._touch_verdicts = {} + return c + + +class TestConnectorTouches: + def test_the_worker_touches_and_records_each_verdict(self): + touched = [] + + def touch(h): + touched.append(h) + return h != 2 + + worker = _k3_worker(SimpleNamespace(codec=SimpleNamespace(touch=touch))) + ids = [StateTouchId(1, 0), StateTouchId(2, 1)] + worker._touch_states(SimpleNamespace(state_touches=ids)) + assert touched == [1, 2] + assert worker._touch_verdicts == {ids[0]: True, ids[1]: False} + + def test_with_no_tier_every_image_is_gone(self): + worker = _k3_worker(None) + worker._touch_states(SimpleNamespace(state_touches=[StateTouchId(1, 0)])) + assert worker._touch_verdicts == {StateTouchId(1, 0): False} + + def test_a_failing_touch_does_not_erase_the_index(self): + def touch(h): + raise RuntimeError("backend hiccup") + + worker = _k3_worker(SimpleNamespace(codec=SimpleNamespace(touch=touch))) + worker._touch_states(SimpleNamespace(state_touches=[StateTouchId(1, 0)])) + assert worker._touch_verdicts == {StateTouchId(1, 0): True} + + def _scheduler(self, monkeypatch): + from atom.kv_transfer.offload.dense.connector import DenseOffloadScheduler + from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( + KimiK3OffloadScheduler, + ) + from atom.kv_transfer.offload.metadata import LMCacheOffloadMetadata + + monkeypatch.setattr( + DenseOffloadScheduler, + "build_connector_meta", + lambda self: LMCacheOffloadMetadata(), + ) + s = KimiK3OffloadScheduler.__new__(KimiK3OffloadScheduler) + s._pending_state_loads = [] + s._pending_state_stores = [] + s._pending_state_touches = [] + s._touch_round = 0 + s._state_touch_verdicts = {} + s._refresh_save_stall = lambda: None + return s + + def test_each_touch_is_a_distinct_identity_drained_once(self, monkeypatch): + s = self._scheduler(monkeypatch) + assert s.enqueue_state_touches([]) is False + assert s.enqueue_state_touches([5, 5]) + meta = s.build_connector_meta() + assert meta.state_touches == [StateTouchId(5, 0), StateTouchId(5, 1)] + assert s.build_connector_meta().state_touches == [] + + def test_one_rank_missing_the_image_makes_it_gone(self, monkeypatch): + from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( + STATE_TOUCH_CHANNEL, + ) + + s = self._scheduler(monkeypatch) + for op, ok in ((StateTouchId(5, 0), True), (StateTouchId(5, 1), False)): + handled = s.connector_completion( + ConnectorCompletion(STATE_TOUCH_CHANNEL, op, ok) + ) + assert handled is None, "handled, but not a terminal save" + s.connector_completion( + ConnectorCompletion(STATE_TOUCH_CHANNEL, StateTouchId(6, 2), True) + ) + assert s.take_state_touch_verdicts() == {5: False, 6: True} + assert s.take_state_touch_verdicts() == {} + + def test_a_negative_round_is_refused(self): + with pytest.raises(ValueError): + StateTouchId(1, -1) + + +# ── the engine: touches published, verdicts applied, admission held ──────── + + +class TestSchedulerPlumbing: + def _sched(self, connector, idx): + from atom.model_engine.scheduler import Scheduler + + s = Scheduler.__new__(Scheduler) + s.kv_connector = connector + s.block_manager = SimpleNamespace(state_offload=idx) + return s + + def test_voted_hashes_reach_the_connector(self, clock): + idx = _index(7) + idx.could_serve(7) + sent = [] + s = self._sched( + SimpleNamespace(enqueue_state_touches=lambda hs: sent.extend(hs) or True), + idx, + ) + s._publish_state_touches() + assert sent == [7] and idx.refresh_hits + + def test_a_connector_without_touches_turns_the_refresh_off(self, clock): + idx = _index(7) + idx.could_serve(7) + s = self._sched(SimpleNamespace(), idx) + s._publish_state_touches() + assert not idx.refresh_hits + clock[0] += so._TIER_VERIFIED_FOR_S + 1 + assert not idx.needs_verify(7), "nothing would ever answer the hold" + + def test_verdicts_are_applied_to_the_index(self, clock): + idx = _index(7, 8) + s = self._sched( + SimpleNamespace(take_state_touch_verdicts=lambda: {7: False, 8: True}), + idx, + ) + s._apply_state_touch_verdicts(idx) + assert not idx.could_serve(7) and idx.could_serve(8) + + +def _paged_bm(): + from test_state_checkpoint import ( + PAGED_COPY_RUNTIME, + ckpt_config, + make_block_manager, + ) + + bm = make_block_manager( + ckpt_config(num_kvcache_blocks=160, pool_entries={"state": 8}), + state_runtime=PAGED_COPY_RUNTIME, + ) + bm.state_offload = StateOffloadIndex(refresh_hits=True) + return bm + + +def _seq(): + from test_state_checkpoint import stateful_seq + + return stateful_seq(list(range(40))) + + +class TestAdmissionHold: + """`record_allocation(defer_unverified=True)` holds an admission that + would load an unconfirmed image, and commits nothing while it does.""" + + def _joint(self, bm, monkeypatch, *, h=77, from_hbm=False): + from atom.model_engine.block_manager import JointBoundaryDecision + + decision = JointBoundaryDecision( + boundary_tokens=16, + boundary_hash=h, + kv_tokens=16, + claim_tokens=16, + state_from_hbm=from_hbm, + ) + commits = [] + monkeypatch.setattr(bm, "_joint_kv_boundary", lambda *a: decision) + monkeypatch.setattr( + bm, "_commit_joint_boundary", lambda seq, d: commits.append(d) + ) + return commits + + def test_an_unconfirmed_joint_image_is_held(self, clock, monkeypatch): + bm = _paged_bm() + commits = self._joint(bm, monkeypatch) + bm.state_offload.note_stored(77) + clock[0] += so._TIER_VERIFIED_FOR_S + 1 + seq = _seq() + + assert not bm.record_allocation(seq, 2, [1, 2], defer_unverified=True) + assert commits == [] and bm.state_offload.loads_deferred_for_verify == 1 + assert bm.state_offload.take_touches() == [77] + + bm.state_offload.apply_touch_verdict(77, True) + assert bm.record_allocation(seq, 2, [1, 2], defer_unverified=True) + assert len(commits) == 1 + + def test_a_gone_image_is_not_parked_on(self, clock, monkeypatch): + """The verdict drops the hash; the retry's boundary search no longer + sees the tier rung (`could_serve` is False), so it plans a lower one.""" + bm = _paged_bm() + self._joint(bm, monkeypatch) + bm.state_offload.note_stored(77) + clock[0] += so._TIER_VERIFIED_FOR_S + 1 + assert not bm.record_allocation(_seq(), 2, [1, 2], defer_unverified=True) + bm.state_offload.take_touches() + bm.state_offload.apply_touch_verdict(77, False) + assert not bm._tier_can_serve(77) + + def test_a_recently_confirmed_image_goes_straight_through(self, clock, monkeypatch): + bm = _paged_bm() + commits = self._joint(bm, monkeypatch) + bm.state_offload.note_stored(77) + assert bm.record_allocation(_seq(), 2, [1, 2], defer_unverified=True) + assert len(commits) == 1 + + def test_a_state_resident_in_hbm_is_never_held(self, clock, monkeypatch): + bm = _paged_bm() + self._joint(bm, monkeypatch, from_hbm=True) + bm.state_offload.note_stored(77) + clock[0] += so._TIER_VERIFIED_FOR_S + 1 + assert bm.record_allocation(_seq(), 2, [1, 2], defer_unverified=True) + + def test_the_fit_probe_path_never_holds(self, clock, monkeypatch): + bm = _paged_bm() + self._joint(bm, monkeypatch) + bm.state_offload.note_stored(77) + clock[0] += so._TIER_VERIFIED_FOR_S + 1 + assert bm.record_allocation(_seq(), 2, [1, 2]) + + def test_an_hbm_hit_holds_only_for_a_tier_sourced_state(self, clock, monkeypatch): + from atom.model_engine.block_manager import JointBoundaryDecision + + bm = _paged_bm() + monkeypatch.setattr( + bm, + "_joint_kv_boundary", + lambda *a: JointBoundaryDecision(skip_reason="no_kv_reach"), + ) + monkeypatch.setattr(bm, "_commit_joint_boundary", lambda seq, d: None) + bm.state_offload.note_stored(2) + clock[0] += so._TIER_VERIFIED_FOR_S + 1 + resident = {2} + monkeypatch.setattr( + bm.paged_state_checkpoints, "contains", lambda h: h in resident + ) + assert bm.record_allocation(_seq(), 2, [1, 2], defer_unverified=True) + resident.clear() + assert not bm.record_allocation(_seq(), 2, [1, 2], defer_unverified=True) + + +def test_the_scheduler_retries_a_held_admission(monkeypatch): + """A held seq stays waiting and is admitted on a later pass.""" + from test_state_checkpoint import ckpt_config, make_scheduler, stateful_seq + + sched = make_scheduler(ckpt_config()) + bm = sched.block_manager + real = bm.record_allocation + holds = [True] + + def record(seq, n, hashes, *, defer_unverified=False): + assert defer_unverified, "admission asks to hold unverified loads" + if holds.pop() if holds else False: + return False + return real(seq, n, hashes, defer_unverified=defer_unverified) + + monkeypatch.setattr(bm, "record_allocation", record) + seq = stateful_seq(list(range(12))) + sched.add(seq) + _, scheduled = sched.schedule() + assert scheduled == {} + assert list(sched.waiting) == [seq] and len(seq.block_table) == 0 + _, scheduled = sched.schedule() + assert seq.id in scheduled and not sched.waiting + + +# ── the startup budget check ─────────────────────────────────────────────── + + +GIB = 1 << 30 + + +class TestBudgetWarning: + def test_a_cap_below_half_the_pool_is_fine(self): + assert state_tier_budget_warning(59 * GIB, 208 * GIB, 112 << 20) is None + + def test_a_cap_the_size_of_the_pool_bounds_nothing(self): + """#2562's accuracy run: a 32 GiB pool with the 32 GiB default cap.""" + msg = state_tier_budget_warning(32 * GIB, 32 * GIB, 112 << 20) + assert "bounds nothing" in msg + assert "LMCACHE_MAX_LOCAL_CPU_SIZE=32.0" in msg + + def test_a_cap_over_half_the_pool_is_called_out(self): + msg = state_tier_budget_warning(40 * GIB, 64 * GIB, 112 << 20) + assert msg is not None and "more than half" in msg + + def test_unknown_sizes_are_not_judged(self): + assert state_tier_budget_warning(32 * GIB, 0, 1) is None + assert state_tier_budget_warning(0, 32 * GIB, 1) is None diff --git a/tests/test_state_tier_reservation.py b/tests/test_state_tier_reservation.py deleted file mode 100644 index 15983b6ac9..0000000000 --- a/tests/test_state_tier_reservation.py +++ /dev/null @@ -1,190 +0,0 @@ -# SPDX-License-Identifier: MIT -"""The state tier's CPU budget is a reservation inside the shared LMCache pool. - -State images and paged KV share one LocalCPU pool. Before, the codec's budget -only bounded the codec's own LRU: KV writes could still push LMCache's LRU over -the images, while the engine-side `StateOffloadIndex` kept advertising them, and -every such hit parked a joint load that then failed and recomputed the whole -prompt. These tests model LocalCPU's eviction rule (only objects whose ref count -is 1 are evictable) and check that a budgeted image survives KV pressure until -the codec itself retires it. -""" - -import dataclasses -import sys -import types -from types import SimpleNamespace - -import pytest - - -@pytest.fixture -def lmcache_key(monkeypatch): - @dataclasses.dataclass(frozen=True) - class CacheEngineKey: - model_name: str - world_size: int - worker_id: int - chunk_hash: int - dtype: object - - utils_module = types.ModuleType("lmcache.utils") - utils_module.CacheEngineKey = CacheEngineKey - lmcache_module = types.ModuleType("lmcache") - lmcache_module.__path__ = [] - lmcache_module.utils = utils_module - monkeypatch.setitem(sys.modules, "lmcache", lmcache_module) - monkeypatch.setitem(sys.modules, "lmcache.utils", utils_module) - return CacheEngineKey - - -class _Obj: - """A MemoryObj's reference count, and nothing else.""" - - def __init__(self, size: int = 1024) -> None: - self.refs = 1 - self.size = size - - def ref_count_up(self) -> None: - self.refs += 1 - - def ref_count_down(self) -> None: - assert self.refs > 0, "reference discharged twice" - self.refs -= 1 - - def get_ref_count(self) -> int: - return self.refs - - def get_size(self) -> int: - return self.size - - -class _LocalCPU: - """`storage_manager` over LocalCPU's rules. - - `batched_put` adopts an object with its own reference unless the key is - already resident, then drops the caller's reference. `evict_lru` is the - pool making room for KV: it skips anything another holder still references. - """ - - def __init__(self, refuse_all: bool = False) -> None: - self.hot: dict = {} - self.refuse_all = refuse_all - self.fail_put = False - - def batched_put(self, keys, objs) -> None: - if self.fail_put: - raise RuntimeError("put failed") - for k, o in zip(keys, objs, strict=True): - if not self.refuse_all and k not in self.hot: - o.ref_count_up() - self.hot[k] = o - o.ref_count_down() - - def get(self, key): - o = self.hot.get(key) - if o is not None: - o.ref_count_up() - return o - - def remove(self, key) -> int: - o = self.hot.pop(key, None) - if o is None: - return 0 - o.ref_count_down() - return 1 - - def evict_lru(self) -> list: - gone = [k for k, o in self.hot.items() if o.get_ref_count() == 1] - for k in gone: - self.hot.pop(k).ref_count_down() - return gone - - -def _codec(storage, max_entries): - from atom.kv_transfer.offload.hybrid.kimi_k3.state_object import StateByteCodec - - codec = StateByteCodec( - backend=SimpleNamespace( - page_unit_views=lambda units: [], state_entry_views=lambda slot: [] - ), - staged=SimpleNamespace(pack=lambda views, obj: None, unpack=lambda o, v: None), - entry_bytes=1024, - model_name="qwen", - world_size=1, - worker_id=0, - layout_id="L", - max_cache_bytes=None if max_entries is None else max_entries * 1024, - ) - codec.bind_storage_manager(storage) - allocated = [] - - def allocate(nbytes): - allocated.append(_Obj(nbytes)) - return allocated[-1] - - codec._allocate = allocate - return codec, allocated - - -def test_kv_pressure_cannot_evict_a_budgeted_image(lmcache_key): - pool = _LocalCPU() - codec, objs = _codec(pool, max_entries=2) - kv = _Obj() - pool.hot["kv-chunk"] = kv # a KV chunk only the pool references - - assert codec.put(1, []) - assert objs[0].refs == 2 # the pool's reference and the codec's - assert pool.evict_lru() == ["kv-chunk"] - assert codec.get(1, 0), "the image the index advertises must still load" - - -def test_the_codec_lru_retires_and_frees_its_oldest_image(lmcache_key): - pool = _LocalCPU() - codec, objs = _codec(pool, max_entries=2) - assert codec.put(1, []) and codec.put(2, []) and codec.put(3, []) - - assert codec.key(1) not in pool.hot - assert objs[0].refs == 0, "a retired image goes back to the allocator" - assert [o.refs for o in objs[1:]] == [2, 2] - assert len(codec._held) == 2 - - -def test_a_resident_image_is_not_stored_twice(lmcache_key): - pool = _LocalCPU() - codec, objs = _codec(pool, max_entries=2) - released = [] - assert codec.put(1, []) - assert codec.put(1, [], on_source_released=lambda: released.append(True)) - - assert len(objs) == 1, "no second image-sized allocation and D2H" - assert released == [True], "the source units must still be handed back" - - -def test_a_refused_put_keeps_no_reference(lmcache_key): - pool = _LocalCPU(refuse_all=True) - codec, objs = _codec(pool, max_entries=2) - assert codec.put(1, []) - - assert objs[0].refs == 0, "holding a refused object would strand its block" - assert codec._held == {} - - -def test_a_failed_put_discharges_both_references(lmcache_key): - pool = _LocalCPU() - pool.fail_put = True - codec, objs = _codec(pool, max_entries=2) - with pytest.raises(RuntimeError): - codec.put(1, []) - - assert objs[0].refs == 0 - assert codec._held == {} - - -def test_an_unbounded_codec_leaves_eviction_to_lmcache(lmcache_key): - pool = _LocalCPU() - codec, objs = _codec(pool, max_entries=None) - assert codec.put(1, []) - - assert objs[0].refs == 1 - assert pool.evict_lru() == [codec.key(1)] From 4d8798f1c8f5a60cace9894335dbe2ffa9190473 Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Tue, 27 Oct 2026 12:55:43 +0000 Subject: [PATCH 03/18] fix(offload): count a state image as confirmed for 60 s, not 5 s An admission that would load a CPU state image waits one scheduling pass for a touch verdict unless the image was confirmed resident recently. With a 5 s window nearly every AgentX load was held (the next turn of a prefix usually comes well after 5 s), and at c48, where the waiting queue is nearly empty, that pass showed up as +0.3 s p50 TTFT on cache-hit requests (first-turn TTFT unchanged). LMCache's LRU can only evict an image after a pool's worth of newer entries; a 208 GiB pool at c64 turns over in about 13 minutes. 60 s stays far inside that and covers the usual store-to-next-turn gap. Co-authored-by: Cursor --- atom/kv_transfer/offload/README.md | 2 +- atom/model_engine/state_offload.py | 9 ++++++--- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index d74a624922..3a3c700bab 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -130,7 +130,7 @@ 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 admission about to load an image not confirmed in the last -few seconds is held for a pass (at most `OFFLOAD_LOOKUP_DEFER_S`) while that +minute 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. `OFFLOAD_STATE_CPU_SIZE` caps how much of the pool images may take; it reserves nothing, and the worker warns at startup diff --git a/atom/model_engine/state_offload.py b/atom/model_engine/state_offload.py index 9a4742d2bc..0e6e347b8a 100644 --- a/atom/model_engine/state_offload.py +++ b/atom/model_engine/state_offload.py @@ -21,9 +21,12 @@ #: How long a hash counts as known-resident after the last evidence that its #: image was in LMCache (stored, loaded, or confirmed by a touch). Every one of -#: those leaves the image at the young end of LMCache's LRU, and the pool turns -#: over in minutes, so this is far inside the time an image needs to age out. -_TIER_VERIFIED_FOR_S = 5.0 +#: those leaves the image at the young end of LMCache's LRU, and LRU can only +#: evict it after a pool's worth of newer entries (a 208 GiB pool at AgentX +#: c64 turns over in about 13 minutes), so this stays well inside the time an +#: image needs to age out. It is also longer than the usual gap between a +#: prefix being stored and its next turn, so most loads are never held. +_TIER_VERIFIED_FOR_S = 60.0 #: A hash voted on every scheduling pass is touched at most this often. The #: touch only has to keep the image near the young end, not track every pass. From eff79dcf97b0c4a96a7da78ce6beed3d70bab23f Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Wed, 28 Oct 2026 02:21:50 +0000 Subject: [PATCH 04/18] fix(offload): bound the verify hold from its start; don't tombstone touches Two issues from review on 4d8798f1c. The admission hold for a touch verdict was bounded by OFFLOAD_LOOKUP_DEFER_S measured from the latest touch send. A held request re-probes every pass, which re-queues the touch, and take_touches re-sends it every 0.5 s, so the bound never ran out: if verdicts stopped arriving (a dropped metadata field, a connector path that does not carry them, a rank that never reports), the request was held forever with nothing logged. The bound now runs from the first hold of the hash (_hold_started_at) and is cleared by any evidence (verdict, store, completed load) or by dropping the hash. When it runs out the admission proceeds as before touches existed, later admissions over that hash are not held again until evidence arrives, the release is counted as state_offload_verify_holds_expired, and the first one logs a warning. KVOutputAggregator tombstoned every connector completion in one bounded ring shared with the state store, source and load-disposition channels. A touch is a new StateTouchId(hash, round) every time, so at c64 touches cycled the ring in under a minute and evicted the tombstones that guard stores and sources against late duplicate reports. round already makes each touch unique, so touches are no longer tombstoned. Co-authored-by: Cursor --- atom/kv_transfer/disaggregation/aggregator.py | 6 +- atom/kv_transfer/disaggregation/types.py | 11 ++-- atom/model_engine/state_offload.py | 63 +++++++++++++++---- tests/test_kv_aggregator.py | 21 +++++++ tests/test_state_tier_lru_refresh.py | 41 +++++++++++- 5 files changed, 121 insertions(+), 21 deletions(-) diff --git a/atom/kv_transfer/disaggregation/aggregator.py b/atom/kv_transfer/disaggregation/aggregator.py index 885b473e24..860b3f3f37 100644 --- a/atom/kv_transfer/disaggregation/aggregator.py +++ b/atom/kv_transfer/disaggregation/aggregator.py @@ -30,6 +30,7 @@ ReqId, SaveCompletionId, SaveOperationId, + StateTouchId, ) logger = logging.getLogger("atom") @@ -155,10 +156,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 diff --git a/atom/kv_transfer/disaggregation/types.py b/atom/kv_transfer/disaggregation/types.py index 9a45b0d4e7..db5ba95561 100644 --- a/atom/kv_transfer/disaggregation/types.py +++ b/atom/kv_transfer/disaggregation/types.py @@ -107,10 +107,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 @@ -129,7 +129,8 @@ class StateTouchId: 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, - because the aggregator tombstones each `(channel, operation_id)`. + so a quorum never merges two touches and the aggregator keeps no tombstone + for them. """ prefix_hash: int diff --git a/atom/model_engine/state_offload.py b/atom/model_engine/state_offload.py index 0e6e347b8a..73302c4ffb 100644 --- a/atom/model_engine/state_offload.py +++ b/atom/model_engine/state_offload.py @@ -118,6 +118,13 @@ def __init__( # the image was resident. self._touch_sent_at: dict[int, float] = {} self._verified_at: dict[int, float] = {} + # hash -> monotonic of the first admission held for its verdict; the + # `defer_s` bound runs from here, so re-sent touches cannot extend it. + # Cleared by any evidence (verdict, store, load) or by dropping `h`. + self._hold_started_at: dict[int, float] = {} + # Hashes whose hold ran out: admissions over them go ahead unheld + # until evidence arrives. + self._hold_expired: set[int] = set() self.touches_sent = 0 # Touches that found the image gone: LMCache evicted it and the index # learned of it before any request parked on it. @@ -125,6 +132,9 @@ def __init__( # Admissions held back a pass because the image they would load had # not been confirmed recently (`needs_verify`). self.loads_deferred_for_verify = 0 + # Holds released because no verdict came back within `defer_s`. Any + # non-zero value means touch verdicts are being lost. + self.verify_holds_expired = 0 # ------------------------------- stores -------------------------------- # def note_stored(self, h: int) -> None: @@ -140,7 +150,7 @@ def note_stored(self, h: int) -> None: # re-stored prefix is treated as freshly used, then trim the cold end. self._hash_lru.pop(h, None) self._hash_lru[h] = None - self._verified_at[h] = monotonic() + self._mark_verified(h) while len(self._hash_lru) > self._hash_cap: old, _ = self._hash_lru.popitem(last=False) self.hashes.discard(old) @@ -157,6 +167,15 @@ def _drop_touch_state(self, h: int) -> None: self._touch_wanted.discard(h) self._touch_sent_at.pop(h, None) self._verified_at.pop(h, None) + self._end_hold(h) + + def _mark_verified(self, h: int) -> None: + self._verified_at[h] = monotonic() + self._end_hold(h) + + def _end_hold(self, h: int) -> None: + self._hold_started_at.pop(h, None) + self._hold_expired.discard(h) def could_serve(self, h: int) -> bool: """Whether a load for `h` could be offered at all: this tier can load @@ -219,7 +238,7 @@ def apply_touch_verdict(self, h: int, resident: bool) -> None: self.touch_misses += 1 self.forget(h) return - self._verified_at[h] = monotonic() + self._mark_verified(h) # The touch made it LMCache's youngest image; keep the cap's order. self._hash_lru.move_to_end(h) @@ -229,6 +248,8 @@ def stop_refreshing(self) -> None: self.refresh_hits = False self._touch_wanted.clear() self._touch_sent_at.clear() + self._hold_started_at.clear() + self._hold_expired.clear() def needs_verify(self, h: int) -> bool: """Whether an admission about to load `h` should wait a pass for a @@ -239,9 +260,14 @@ def needs_verify(self, h: int) -> bool: the verdict is normally in: resident lets the load go ahead, gone has already dropped `h`, so the admission picks a lower boundary (an HBM checkpoint, an older CPU image, or a plain prefill) instead of parking - and recomputing everything after the miss. A verdict that has not come - back within `defer_s` of its touch stops holding the admission, which - then proceeds as it did before touches existed. + and recomputing everything after the miss. + + The hold is bounded by `defer_s` from the first time `h` was held, not + from its latest touch: a held request re-probes every pass and its + touch is re-sent every `_TIER_TOUCH_INTERVAL_S`, so a bound measured + from the send would never run out if verdicts stopped arriving. Once + it runs out, admissions over `h` proceed as they did before touches + existed until any evidence about `h` arrives. """ if not self.refresh_hits or h not in self.hashes: return False @@ -252,12 +278,24 @@ def needs_verify(self, h: int) -> bool: verified = self._verified_at.get(h) if verified is not None and now - verified <= _TIER_VERIFIED_FOR_S: return False - sent = self._touch_sent_at.get(h) - if sent is None or (verified is not None and verified >= sent): - # No touch outstanding: ask for one; it leaves this pass. - self._touch_wanted.add(h) - return True - return now - sent <= defer_s + if h in self._hold_expired: + return False + started = self._hold_started_at.setdefault(h, now) + if now - started > defer_s: + self._hold_expired.add(h) + self.verify_holds_expired += 1 + if self.verify_holds_expired == 1: + logger.warning( + "state offload: no touch verdict for a CPU state image " + "within OFFLOAD_LOOKUP_DEFER_S=%.1fs; admitting without it. " + "Touch verdicts are being lost (counted as " + "state_offload_verify_holds_expired).", + defer_s, + ) + return False + # Ask for a touch; `take_touches` throttles the resend. + self._touch_wanted.add(h) + return True # -------------------------------- loads -------------------------------- # def request_load(self, req_id, h: int) -> bool: @@ -300,7 +338,7 @@ def complete_load(self, req_id) -> None: # order so the cap trims what the codec will trim. if h in self._hash_lru: self._hash_lru.move_to_end(h) - self._verified_at[h] = monotonic() + self._mark_verified(h) def fail_load(self, req_id) -> None: """No bytes came back. Retract the claim as well as counting it. @@ -358,6 +396,7 @@ def stats(self) -> dict[str, int]: "touches_sent": self.touches_sent, "touch_misses": self.touch_misses, "loads_deferred_for_verify": self.loads_deferred_for_verify, + "verify_holds_expired": self.verify_holds_expired, } diff --git a/tests/test_kv_aggregator.py b/tests/test_kv_aggregator.py index b3e7cbfb52..c36926703b 100644 --- a/tests/test_kv_aggregator.py +++ b/tests/test_kv_aggregator.py @@ -15,6 +15,8 @@ from atom.kv_transfer.disaggregation.types import ( LoadOperationId, SaveOperationId, + StateStoreOperationId, + StateTouchId, ) from atom.kv_transfer.offload.hybrid.dsv4.connector import ( DSV4_CHECKPOINT_SAVE_CHANNEL, @@ -413,6 +415,25 @@ def test_terminal_tombstones_are_bounded_and_reset(self): agg.reset() assert agg.terminal_tombstone_count == (0, 0) + def test_state_touches_do_not_evict_store_and_source_tombstones(self): + agg = KVOutputAggregator(world_size=2, terminal_tombstone_limit=4) + store = _completion("k3_state_index", StateStoreOperationId(9, 0)) + source = _completion("k3_state_source", StateStoreOperationId(9, 0)) + both = KVConnectorOutput(connector_completions={store, source}) + assert agg.aggregate([both, both]).connector_completions == {store, source} + + for round_ in range(100): + touch = KVConnectorOutput( + connector_completions={ + _completion("k3_state_touch", StateTouchId(9, round_)) + } + ) + assert agg.aggregate([touch, touch]).connector_completions + assert agg.terminal_tombstone_count == (0, 2) + + late = agg.aggregate([both, both]) + assert late.connector_completions == set(), "late duplicates still dropped" + def test_legacy_raw_save_ids_are_not_tombstoned_for_reuse(self): agg = KVOutputAggregator(world_size=2) terminal = [ diff --git a/tests/test_state_tier_lru_refresh.py b/tests/test_state_tier_lru_refresh.py index 709da83582..2190a7d017 100644 --- a/tests/test_state_tier_lru_refresh.py +++ b/tests/test_state_tier_lru_refresh.py @@ -368,15 +368,50 @@ def test_a_gone_verdict_ends_the_hold_by_dropping_the_hash(self, clock): def test_a_lost_verdict_holds_no_longer_than_the_lookup_defer( self, clock, monkeypatch ): + monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "2") + idx = _index(7) + clock[0] += so._TIER_VERIFIED_FOR_S + 1 + # A held request re-probes every pass and its touch is re-sent every + # interval; verdicts never come back. Re-sends must not extend the hold. + held_for, step = 0.0, 0.05 + while idx.needs_verify(7): + assert held_for <= 2 + 1e-6, "the hold outlived OFFLOAD_LOOKUP_DEFER_S" + assert idx.could_serve(7) + idx.take_touches() + clock[0] += step + held_for += step + assert held_for >= 2 + assert idx.touches_sent > 1, "the touch was re-sent while held" + assert idx.verify_holds_expired == 1 + assert idx.stats()["verify_holds_expired"] == 1 + assert not idx.needs_verify(7), "later admissions are not held again" + assert idx.verify_holds_expired == 1 + + def test_evidence_after_an_expired_hold_rearms_it(self, clock, monkeypatch): + monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "2") + idx = _index(7) + clock[0] += so._TIER_VERIFIED_FOR_S + 1 + assert idx.needs_verify(7) + clock[0] += 3 + assert not idx.needs_verify(7) + idx.apply_touch_verdict(7, True) + clock[0] += so._TIER_VERIFIED_FOR_S + 1 + assert idx.needs_verify(7), "a fresh hold, timed from now" + clock[0] += 1 + assert idx.needs_verify(7) + + def test_a_verdict_ends_the_hold_timer(self, clock, monkeypatch): monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "2") idx = _index(7) clock[0] += so._TIER_VERIFIED_FOR_S + 1 assert idx.needs_verify(7) - idx.take_touches() clock[0] += 1.5 + idx.apply_touch_verdict(7, True) + clock[0] += so._TIER_VERIFIED_FOR_S + 1 assert idx.needs_verify(7) - clock[0] += 1.0 - assert not idx.needs_verify(7) + clock[0] += 1.5 + assert idx.needs_verify(7), "the earlier 1.5 s does not count" + assert idx.verify_holds_expired == 0 def test_a_zero_defer_never_holds(self, clock, monkeypatch): monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "0") From c54eda26a383a2997e0ca54e3d5872e3e5a39b42 Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Fri, 30 Oct 2026 03:09:04 +0000 Subject: [PATCH 05/18] refactor(offload): share the resume-hash pick, size the verify window from the pool, keep could_serve pure Review round 3 on #2668 (ZhangLirong-amd): 1. `allocate` and the admission hold (`_tier_state_unverified`) now pick the state hash through one helper, `BlockManager._state_resume_hash`: the joint boundary when there is one, none for a replay, else the HBM hit. Behaviour is unchanged (`transfer.replays` never coexists with PAGE checkpoints, and `_attach_state_slots` already dropped the hash for it). 2. The confirmation window is no longer a hard-coded 60 s. It is derived from the LMCache CPU pool (`LMCACHE_MAX_LOCAL_CPU_SIZE`, read with the chunk size in `state_tier_lmcache_sizes`): a quarter of the pool's turnover at an assumed 1 GiB/s fill ceiling (4x the AgentX c64 rate), i.e. 0.25 s per GiB, clamped to [5 s, 60 s]; 52 s for the 208.6 GiB AgentX pool, the floor when the pool is unknown. No new environment variable. 3. `StateOffloadIndex.could_serve` is a pure check again. Voters register the hits they count with the new `want_touch`: the resumable scan (`PagedStateCheckpointCoordinator._reachable`) and an offered load (`BlockManager._request_state_load`). `_tier_can_serve` stays pure. 4. `OFFLOAD_LOOKUP_DEFER_S` is read once when the index is built (only when it refreshes hits) instead of on every `needs_verify`; tests pass `defer_s`. Tests: `tests/test_state_tier_lru_refresh.py` (63, +16): could_serve queues nothing, voters register touches, the defer is read once, the window's formula and bounds and the block manager feeding it the pool size, the resume-hash helper, and an end-to-end check that the hold asks about the same hash `allocate` then loads. --- atom/kv_transfer/offload/README.md | 10 +- atom/model_engine/block_manager.py | 86 +++++--- atom/model_engine/page_unit_checkpoint.py | 8 +- atom/model_engine/state_offload.py | 116 +++++++--- tests/test_page_unit_checkpoint.py | 3 + tests/test_state_checkpoint.py | 6 + tests/test_state_tier_lru_refresh.py | 248 ++++++++++++++++++---- 7 files changed, 372 insertions(+), 105 deletions(-) diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index 3a3c700bab..4d7e759e7c 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -129,10 +129,12 @@ every hit the index votes for is also sent to the workers as a *touch* 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 admission about to load an image not confirmed in the last -minute 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. `OFFLOAD_STATE_CPU_SIZE` caps how much of +verdict lands. 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`. diff --git a/atom/model_engine/block_manager.py b/atom/model_engine/block_manager.py index f056b51c61..810b5e496e 100644 --- a/atom/model_engine/block_manager.py +++ b/atom/model_engine/block_manager.py @@ -335,8 +335,8 @@ def __init__( # releases holds a whole image out of the pool forever. from atom.model_engine.state_offload import ( state_tier_capability, - state_tier_chunk_tokens, state_tier_cpu_bytes, + state_tier_lmcache_sizes, ) # A capability derived from the whole config, not the connector's name. @@ -461,13 +461,16 @@ def __init__( # The LMCache chunk size in tokens, read where the config is rather than # off the connector object. self._joint_chunk_tokens = 0 + lmcache_pool_bytes = 0 # Only when a joint KV load is possible at all. Reading it reaches into # LMCache, an optional dependency: probing on an engine that hosts no # state tier logged a full `ModuleNotFoundError` traceback at WARNING # on every start, for a number that engine has no use for. The # capability above is the same gate the tier itself uses. if kv_offload_enabled: - self._joint_chunk_tokens = state_tier_chunk_tokens(config) + self._joint_chunk_tokens, lmcache_pool_bytes = state_tier_lmcache_sizes( + config + ) self.state_offload: StateOffloadIndex | None = None # (req_id, hash, target_group) admitted this pass and not yet handed to # the connector. Kept here rather than in the index because the slot @@ -518,6 +521,8 @@ def __init__( # Off again (`stop_refreshing`) if the connector turns out not # to carry state touches. refresh_hits=True, + # Sizes how long a confirmed image is trusted. + pool_bytes=lmcache_pool_bytes, ) # Attached rather than passed at construction: the coordinator is # built before the switch is read (it needs `checkpoint_spec`, @@ -1013,7 +1018,7 @@ def _joint_kv_boundary( if not seq.has_per_req_cache or self.paged_state_checkpoints is None: return self._no_joint("no_paged_checkpoints") hbs = self._hash_block_size() - # Resolved once from config at construction (`state_tier_chunk_tokens`), + # Resolved once from config at construction (`state_tier_lmcache_sizes`), # not off the connector object: the scheduler holds whatever # `get_kvconnector` returned, and one without `chunk_size` would zero # this and disable the feature silently. @@ -1579,25 +1584,27 @@ def _tier_state_unverified( """Whether `allocate` would load this seq's state from the CPU tier under a hash not confirmed resident recently. - Mirrors which hash `allocate` hands `_attach_state_slots`: the joint - boundary when there is one, else the HBM hit's last block; and the - tier only serves it when HBM does not hold it. + Judges the hash `allocate` will hand `_attach_state_slots` + (`_state_resume_hash`), which the tier only serves when HBM does not + hold it. """ offload = self.state_offload if offload is None or not offload.refresh_hits or not seq.has_per_req_cache: return False - if decision.skip_reason is None: - if decision.state_from_hbm: - return False - h = decision.boundary_hash - else: - if not 0 < num_cached_blocks <= len(block_hashes): - return False - if self.state.transfer.replays or ( - self.state_replay_tokens and seq.replays_hit - ): - return False # a replay rebuilds the state; nothing is loaded - h = block_hashes[num_cached_blocks - 1] + joint = decision.skip_reason is None + if joint and decision.state_from_hbm: + return False + hit_hash = ( + block_hashes[num_cached_blocks - 1] + if 0 < num_cached_blocks <= len(block_hashes) + else -1 + ) + h = self._state_resume_hash( + seq, decision.boundary_hash if joint else -1, hit_hash + ) + if h == -1: + return False + if not joint: if self.paged_state_checkpoints is not None: if self.paged_state_checkpoints.contains(h): return False @@ -1605,6 +1612,23 @@ def _tier_state_unverified( return False return offload.needs_verify(h) + def _state_resume_hash(self, seq: Sequence, joint_hash: int, hit_hash: int) -> int: + """The hash whose state `seq` resumes from, or -1 for none. + + The joint boundary when there is one (above the HBM hit by + construction), else the HBM hit's last block. A replay rebuilds the + state instead of restoring it, so it resumes from none. `allocate` + hands this to `_attach_state_slots`, and `_tier_state_unverified` asks + about it, so the admission hold judges the hash that will be loaded. + """ + if joint_hash != -1: + return joint_hash + if self.state.transfer.replays or ( + self.state_replay_tokens and seq.replays_hit + ): + return -1 + return hit_hash + def retreat_joint_to_hbm(self, seq: Sequence) -> bool: """Drop a failed unaligned joint and resume at the HBM hit it lowered. @@ -1800,13 +1824,16 @@ def allocate(self, seq: Sequence, num_cached_blocks: int = 0) -> bool: # last one actually claimed, whichever way the loop exits. if i < num_cached_blocks: hit_hash = h - if self.state_replay_tokens and seq.replays_hit: - # A replay starts from a fresh state slot, never a checkpoint: it - # rebuilds the state below whatever the claim above reached. - hit_hash = -1 - if num_cached_blocks < requested_blocks: - # A KV-only load past the hit would start above a gap. - seq.replay_kv_load = False + # A replay starts from a fresh state slot, never a checkpoint: it + # rebuilds the state below whatever the claim above reached + # (`_state_resume_hash` gives it none). A KV-only load past the hit + # would start above a gap. + if ( + self.state_replay_tokens + and seq.replays_hit + and num_cached_blocks < requested_blocks + ): + seq.replay_kv_load = False # Pin the restore before fresh blocks can evict its checkpoint. state_holds = True if seq.has_per_req_cache and self.paged_state_checkpoints is not None: @@ -1822,9 +1849,9 @@ def allocate(self, seq: Sequence, num_cached_blocks: int = 0) -> bool: # joint boundary and `_claim_after_load` raised `num_cached_tokens` # to it, so the forward resumed over a prefix the state does not # cover. Silent wrong output, no exception. - joint_hash = seq.offload_joint.boundary_hash state_holds = self._attach_state_slots( - seq, joint_hash if joint_hash != -1 else hit_hash + seq, + self._state_resume_hash(seq, seq.offload_joint.boundary_hash, hit_hash), ) for _ in range(len(seq.block_table), self.num_pool_blocks(len(seq))): seq.block_table.append(self._fresh_block()) @@ -1841,9 +1868,9 @@ def allocate(self, seq: Sequence, num_cached_blocks: int = 0) -> bool: # above the HBM hit by construction. `num_cached_tokens` stays at # the HBM prefix until the KV leg lands, so the forward covers # `[hbm, num_prompt)` if anything goes wrong from here. - joint_hash = seq.offload_joint.boundary_hash state_holds = self._attach_state_slots( - seq, joint_hash if joint_hash != -1 else hit_hash + seq, + self._state_resume_hash(seq, seq.offload_joint.boundary_hash, hit_hash), ) if seq.has_per_req_cache: seq._state_initialized_after_alloc = False @@ -2091,6 +2118,7 @@ def _request_state_load(self, seq: Sequence, hit_hash: int) -> bool: return False if not self.state_offload.request_load(seq.id, hit_hash): return False + self.state_offload.want_touch(hit_hash) seq.offload_joint.load_hash = hit_hash self._state_loads.append((seq.id, hit_hash, seq.state_slot)) return True diff --git a/atom/model_engine/page_unit_checkpoint.py b/atom/model_engine/page_unit_checkpoint.py index 44406a976b..2adb966c81 100644 --- a/atom/model_engine/page_unit_checkpoint.py +++ b/atom/model_engine/page_unit_checkpoint.py @@ -1072,10 +1072,16 @@ def _reachable(self, h: int) -> bool: `hashes` from its own stores but cannot load them back, and voting a hit it would then refuse strands the scan at the tier rung -- skipping a still-resident HBM rung -- and forfeits the checkpoint to a recompute. + + A tier hit accepted here is a vote, so it is registered for a touch + (`want_touch`). """ if self.store.contains(h): return True - return self.offload is not None and self.offload.could_serve(h) + if self.offload is None or not self.offload.could_serve(h): + return False + self.offload.want_touch(h) + return True def checkpoint(self, seq: Sequence, boundary_blocks: int, h: int) -> None: """File a boundary to be stored, keyed by hash rather than by seq. diff --git a/atom/model_engine/state_offload.py b/atom/model_engine/state_offload.py index 73302c4ffb..36631ff488 100644 --- a/atom/model_engine/state_offload.py +++ b/atom/model_engine/state_offload.py @@ -19,14 +19,15 @@ logger = logging.getLogger("atom") -#: How long a hash counts as known-resident after the last evidence that its -#: image was in LMCache (stored, loaded, or confirmed by a touch). Every one of -#: those leaves the image at the young end of LMCache's LRU, and LRU can only -#: evict it after a pool's worth of newer entries (a 208 GiB pool at AgentX -#: c64 turns over in about 13 minutes), so this stays well inside the time an -#: image needs to age out. It is also longer than the usual gap between a -#: prefix being stored and its next turn, so most loads are never held. -_TIER_VERIFIED_FOR_S = 60.0 +#: The fastest one worker's LMCache pool is assumed to take in new entries. +#: AgentX c64 on one GPU stored about 0.25 GiB/s (970 GB/h); four times that +#: leaves room for a prefill-heavier load or a model with more KV per token. +_POOL_FILL_GIB_PER_S = 1.0 +#: The share of the pool's turnover at that rate an image is trusted for. +_POOL_TRUSTED_FRACTION = 0.25 +#: Bounds on `tier_verified_for_s`. See there. +_TIER_VERIFIED_MIN_S = 5.0 +_TIER_VERIFIED_MAX_S = 60.0 #: A hash voted on every scheduling pass is touched at most this often. The #: touch only has to keep the image near the young end, not track every pass. @@ -40,8 +41,8 @@ class StateOffloadIndex: can drop bytes at any time, so on its own a hash here means "was stored once", never "is still there". - With `refresh_hits`, every hit the index votes for is also *touched*: the - hash is queued for `take_touches`, the worker refreshes the image's + With `refresh_hits`, every hit a voter registers (`want_touch`) is also + *touched*: the hash is queued for `take_touches`, the worker refreshes the image's position in LMCache's LRU the same way a KV lookup refreshes its chunks, and reports whether the image was still there (`apply_touch_verdict`). That keeps a prefix's state image as warm as its KV, and lets a hash LMCache @@ -57,6 +58,8 @@ def __init__( can_load: bool = True, max_cpu_entries: int | None = None, refresh_hits: bool = False, + pool_bytes: int = 0, + defer_s: float | None = None, ) -> None: # The two legs are separately granted. A `kv_producer` role saves and # never loads, a `kv_consumer` loads and never saves, and offering a @@ -112,6 +115,12 @@ def __init__( # connector carries no touches (the vLLM plugin), where the queue would # only grow. self.refresh_hits = bool(refresh_hits) + # Both fixed for the server's lifetime; `needs_verify` runs on the + # admission path. `pool_bytes` is one worker's LMCache CPU pool. + self.verified_for_s = tier_verified_for_s(pool_bytes) + if defer_s is None: + defer_s = verify_defer_s() if self.refresh_hits else 0.0 + self.defer_s = float(defer_s) # Hashes voted for since the last `take_touches`. self._touch_wanted: set[int] = set() # hash -> monotonic of the last touch handed out / the last evidence @@ -179,13 +188,8 @@ def _end_hold(self, h: int) -> None: def could_serve(self, h: int) -> bool: """Whether a load for `h` could be offered at all: this tier can load - AND believes it holds `h`. - - A True answer is a hit vote, so under `refresh_hits` it also queues `h` - for a touch. This is the one predicate every voter goes through, which - is what makes it the place to refresh: a prefix whose state is voted - for on every pass stays as warm in LMCache as its KV chunks, which the - KV lookup refreshes on every pass. + AND believes it holds `h`. Pure; a voter that acts on a True answer + registers the hit with `want_touch`. The membership+capability half of `request_load`'s guard (minus its per-request in-flight check), factored out so the admission-path voters @@ -198,13 +202,21 @@ def could_serve(self, h: int) -> bool: still-resident HBM rung, and `request_load` returned False -- the boundary disowned and the HBM checkpoint forfeited to a full recompute. """ - if not (self.can_load and h in self.hashes): - return False - if self.refresh_hits: - self._touch_wanted.add(h) - return True + return self.can_load and h in self.hashes # ------------------------------- touches ------------------------------- # + def want_touch(self, h: int) -> None: + """A voter counted `h` as a hit: under `refresh_hits`, queue it for + `take_touches`. + + Called by every voter on the hits it votes for (the resumable scan and + `BlockManager._request_state_load`), so a prefix whose state is voted + for on every pass stays as warm in LMCache as its KV chunks, which the + KV lookup refreshes on every pass. + """ + if self.refresh_hits and h in self.hashes: + self._touch_wanted.add(h) + def take_touches(self) -> list[int]: """Hashes to touch in LMCache this step, drained. @@ -256,27 +268,26 @@ def needs_verify(self, h: int) -> bool: touch verdict instead of parking on it now. True while `h` has not been confirmed resident within - `_TIER_VERIFIED_FOR_S` and its touch is still pending. A pass later + `verified_for_s` and its touch is still pending. A pass later the verdict is normally in: resident lets the load go ahead, gone has already dropped `h`, so the admission picks a lower boundary (an HBM checkpoint, an older CPU image, or a plain prefill) instead of parking and recomputing everything after the miss. - The hold is bounded by `defer_s` from the first time `h` was held, not + The hold is bounded by `defer_s` (`OFFLOAD_LOOKUP_DEFER_S`, read once + at construction) from the first time `h` was held, not from its latest touch: a held request re-probes every pass and its touch is re-sent every `_TIER_TOUCH_INTERVAL_S`, so a bound measured from the send would never run out if verdicts stopped arriving. Once it runs out, admissions over `h` proceed as they did before touches existed until any evidence about `h` arrives. """ - if not self.refresh_hits or h not in self.hashes: - return False - defer_s = verify_defer_s() - if defer_s <= 0: + defer_s = self.defer_s + if not self.refresh_hits or defer_s <= 0 or h not in self.hashes: return False now = monotonic() verified = self._verified_at.get(h) - if verified is not None and now - verified <= _TIER_VERIFIED_FOR_S: + if verified is not None and now - verified <= self.verified_for_s: return False if h in self._hold_expired: return False @@ -402,7 +413,7 @@ def stats(self) -> dict[str, int]: def verify_defer_s() -> float: """How long an admission may wait for a touch verdict - (`StateOffloadIndex.needs_verify`). + (`StateOffloadIndex.needs_verify`), read once per index. `OFFLOAD_LOOKUP_DEFER_S`: the same question it already answers for a KV lookup in flight -- how long admission may pass over a request whose tier @@ -413,6 +424,36 @@ def verify_defer_s() -> float: return float(envs.OFFLOAD_LOOKUP_DEFER_S) +def tier_verified_for_s(pool_bytes: int) -> float: + """How long a hash counts as known-resident after the last evidence that + its image was in LMCache (stored, loaded, or confirmed by a touch). + + Each of those leaves the image at the young end of LMCache's LRU, which + evicts it only after a whole pool of newer entries has come in: the safe + window is the pool's turnover, `pool_bytes / fill rate`. The fill rate is + the workload's and not known at startup, so this takes a ceiling, + `_POOL_FILL_GIB_PER_S`, and trusts `_POOL_TRUSTED_FRACTION` of the + turnover at that rate: 0.25 s per GiB, 52 s for the 208.6 GiB AgentX pool, + whose turnover at c64 was about 14 minutes. + + Clamped to [`_TIER_VERIFIED_MIN_S`, `_TIER_VERIFIED_MAX_S`] = [5 s, 60 s]. + The floor (pools under 20 GiB) keeps the window well above a touch round + trip, `_TIER_TOUCH_INTERVAL_S` plus a step, or nearly every load would be + held. The ceiling (pools over 240 GiB) is the largest window measured + (AgentX c48 / c64); past it a fill rate above the assumed one would let + images evicted minutes earlier through. An unknown pool (0) gets the floor. + + Erring either way is a cost, not wrong output: too long admits a load that + may miss and recompute, which is what every load did before touches; too + short holds a load a pass for a verdict it did not need. + """ + if pool_bytes <= 0: + return _TIER_VERIFIED_MIN_S + turnover_s = pool_bytes / float(1 << 30) / _POOL_FILL_GIB_PER_S + window = turnover_s * _POOL_TRUSTED_FRACTION + return min(max(window, _TIER_VERIFIED_MIN_S), _TIER_VERIFIED_MAX_S) + + #: Offload layouts whose worker half builds a `StateOffloadTier`. Not every #: `lmcache_offload` does: `dense` has no per-request state at all and `hybrid` #: (DSV4) keeps its own in the SLOT sidecar. `qwen` subclasses the K3 tier. @@ -622,8 +663,10 @@ def state_tier_budget_warning( return None -def state_tier_chunk_tokens(config) -> int: - """The LMCache chunk size in tokens for the joint KV leg, or 0. +def state_tier_lmcache_sizes(config) -> tuple[int, int]: + """The LMCache chunk size in tokens for the joint KV leg, and one worker's + CPU pool in bytes (`LMCACHE_MAX_LOCAL_CPU_SIZE`) for `tier_verified_for_s`; + (0, 0) when LMCache's config cannot be read. Decided from config for the same reason `state_tier_capability` is: this runs in the engine process at `BlockManager.__init__`, and the scheduler @@ -652,7 +695,9 @@ def state_tier_chunk_tokens(config) -> int: sub, _why = _offload_subconfig(kvcfg) if sub is not None: kvcfg = sub - return int(build_lmcache_config(kvcfg).chunk_size) + cfg = build_lmcache_config(kvcfg) + pool_gib = float(getattr(cfg, "max_local_cpu_size", 0) or 0) + return int(cfg.chunk_size), int(pool_gib * (1 << 30)) except Exception: # Blind on purpose: this runs at model load, the import reaches a # third-party package that may be absent entirely, and the only @@ -660,10 +705,11 @@ def state_tier_chunk_tokens(config) -> int: # stays off. Refusing to start would be worse. logger.warning( "state offload: could not read the LMCache chunk size; the joint " - "KV load needs it and stays off", + "KV load needs it and stays off, and the CPU pool size is taken " + "as unknown", exc_info=True, ) - return 0 + return 0, 0 #: The connector backends whose worker half can build a `StateOffloadTier`. diff --git a/tests/test_page_unit_checkpoint.py b/tests/test_page_unit_checkpoint.py index 5e23f4408a..a196127412 100644 --- a/tests/test_page_unit_checkpoint.py +++ b/tests/test_page_unit_checkpoint.py @@ -646,6 +646,9 @@ def __init__(self, *hashes): def could_serve(self, h): return h in self.hashes + def want_touch(self, h): + pass + @staticmethod def coordinator(num_units=40, offload=None): from atom.model_engine.page_unit_checkpoint import ( diff --git a/tests/test_state_checkpoint.py b/tests/test_state_checkpoint.py index a480163131..47bc047caa 100644 --- a/tests/test_state_checkpoint.py +++ b/tests/test_state_checkpoint.py @@ -1405,6 +1405,9 @@ def __init__(self, *hashes): def could_serve(self, h): return h in self.hashes + def want_touch(self, h): + pass + def request_load(self, req_id, h): if h not in self.hashes: return False @@ -4493,6 +4496,9 @@ def __init__(self, *hashes): def could_serve(self, h): return self.can_load and h in self.hashes + def want_touch(self, h): + pass + def request_load(self, req_id, h): if h not in self.hashes: return False diff --git a/tests/test_state_tier_lru_refresh.py b/tests/test_state_tier_lru_refresh.py index 022f7f65b0..9c915fb4fa 100644 --- a/tests/test_state_tier_lru_refresh.py +++ b/tests/test_state_tier_lru_refresh.py @@ -285,34 +285,45 @@ def test_without_a_local_cpu_backend_a_touch_asks_presence(self, lmcache_key): # ── the index: votes queue touches, verdicts keep it in step ─────────────── -def _index(*hashes, refresh=True): - idx = StateOffloadIndex(refresh_hits=refresh) +#: Past every confirmation window `tier_verified_for_s` can give. +STALE = so._TIER_VERIFIED_MAX_S + 1 + + +def _index(*hashes, refresh=True, **kw): + idx = StateOffloadIndex(refresh_hits=refresh, **kw) for h in hashes: idx.note_stored(h) return idx class TestIndexTouches: - def test_a_hit_vote_queues_a_touch(self, clock): + def test_could_serve_queues_nothing(self, clock): idx = _index(7) assert idx.could_serve(7) and not idx.could_serve(8) + assert idx.take_touches() == [] + + def test_a_registered_hit_queues_a_touch(self, clock): + idx = _index(7) + idx.want_touch(7) + idx.want_touch(8) # not indexed: nothing to touch assert idx.take_touches() == [7] assert idx.take_touches() == [] def test_no_touches_without_refresh(self, clock): idx = _index(7, refresh=False) assert idx.could_serve(7) + idx.want_touch(7) assert idx.take_touches() == [] def test_a_hash_voted_every_pass_is_touched_at_an_interval(self, clock): idx = _index(7) - idx.could_serve(7) + idx.want_touch(7) assert idx.take_touches() == [7] clock[0] += so._TIER_TOUCH_INTERVAL_S / 2 - idx.could_serve(7) + idx.want_touch(7) assert idx.take_touches() == [] clock[0] += so._TIER_TOUCH_INTERVAL_S - idx.could_serve(7) + idx.want_touch(7) assert idx.take_touches() == [7] assert idx.touches_sent == 2 @@ -339,7 +350,7 @@ def test_a_verdict_for_a_forgotten_hash_is_ignored(self, clock): def test_stop_refreshing_drops_pending_touches(self, clock): idx = _index(7) - idx.could_serve(7) + idx.want_touch(7) idx.stop_refreshing() assert idx.take_touches() == [] assert not idx.needs_verify(7) @@ -351,7 +362,7 @@ def test_a_fresh_store_needs_no_verdict(self, clock): def test_a_stale_hash_waits_for_its_touch(self, clock): idx = _index(7) - clock[0] += so._TIER_VERIFIED_FOR_S + 1 + clock[0] += STALE assert idx.needs_verify(7) assert idx.take_touches() == [7], "the hold asks for its own touch" assert idx.needs_verify(7), "still out" @@ -360,18 +371,15 @@ def test_a_stale_hash_waits_for_its_touch(self, clock): def test_a_gone_verdict_ends_the_hold_by_dropping_the_hash(self, clock): idx = _index(7) - clock[0] += so._TIER_VERIFIED_FOR_S + 1 + clock[0] += STALE assert idx.needs_verify(7) idx.take_touches() idx.apply_touch_verdict(7, False) assert not idx.needs_verify(7) and not idx.could_serve(7) - def test_a_lost_verdict_holds_no_longer_than_the_lookup_defer( - self, clock, monkeypatch - ): - monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "2") - idx = _index(7) - clock[0] += so._TIER_VERIFIED_FOR_S + 1 + def test_a_lost_verdict_holds_no_longer_than_the_lookup_defer(self, clock): + idx = _index(7, defer_s=2) + clock[0] += STALE # A held request re-probes every pass and its touch is re-sent every # interval; verdicts never come back. Re-sends must not extend the hold. held_for, step = 0.0, 0.05 @@ -388,27 +396,25 @@ def test_a_lost_verdict_holds_no_longer_than_the_lookup_defer( assert not idx.needs_verify(7), "later admissions are not held again" assert idx.verify_holds_expired == 1 - def test_evidence_after_an_expired_hold_rearms_it(self, clock, monkeypatch): - monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "2") - idx = _index(7) - clock[0] += so._TIER_VERIFIED_FOR_S + 1 + def test_evidence_after_an_expired_hold_rearms_it(self, clock): + idx = _index(7, defer_s=2) + clock[0] += STALE assert idx.needs_verify(7) clock[0] += 3 assert not idx.needs_verify(7) idx.apply_touch_verdict(7, True) - clock[0] += so._TIER_VERIFIED_FOR_S + 1 + clock[0] += STALE assert idx.needs_verify(7), "a fresh hold, timed from now" clock[0] += 1 assert idx.needs_verify(7) - def test_a_verdict_ends_the_hold_timer(self, clock, monkeypatch): - monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "2") - idx = _index(7) - clock[0] += so._TIER_VERIFIED_FOR_S + 1 + def test_a_verdict_ends_the_hold_timer(self, clock): + idx = _index(7, defer_s=2) + clock[0] += STALE assert idx.needs_verify(7) clock[0] += 1.5 idx.apply_touch_verdict(7, True) - clock[0] += so._TIER_VERIFIED_FOR_S + 1 + clock[0] += STALE assert idx.needs_verify(7) clock[0] += 1.5 assert idx.needs_verify(7), "the earlier 1.5 s does not count" @@ -417,17 +423,81 @@ def test_a_verdict_ends_the_hold_timer(self, clock, monkeypatch): def test_a_zero_defer_never_holds(self, clock, monkeypatch): monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "0") idx = _index(7) - clock[0] += so._TIER_VERIFIED_FOR_S + 1 + assert idx.defer_s == 0 + clock[0] += STALE assert not idx.needs_verify(7) + def test_the_lookup_defer_is_read_once_at_construction(self, clock, monkeypatch): + """`needs_verify` is on the admission path; it must not re-read and + re-parse `OFFLOAD_LOOKUP_DEFER_S` on every call.""" + monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "2") + idx = _index(7) + assert idx.defer_s == 2.0 + + def unread(): + raise AssertionError("OFFLOAD_LOOKUP_DEFER_S read on the admission path") + + monkeypatch.setattr(so, "verify_defer_s", unread) + monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "0") + clock[0] += STALE + assert idx.needs_verify(7), "the value read at construction still holds" + + def test_an_index_without_refresh_never_reads_the_defer(self, monkeypatch): + def unread(): + raise AssertionError("read for an index that never holds") + + monkeypatch.setattr(so, "verify_defer_s", unread) + assert StateOffloadIndex().defer_s == 0 + def test_a_completed_load_counts_as_verified(self, clock): idx = _index(7) - clock[0] += so._TIER_VERIFIED_FOR_S + 1 + clock[0] += STALE assert idx.request_load("r", 7) idx.complete_load("r") assert not idx.needs_verify(7) +class TestVerifiedWindow: + """How long a confirmed image is trusted scales with the LMCache pool it + has to age out of, within [5 s, 60 s].""" + + GIB = 1 << 30 + + def test_the_agentx_pool_keeps_about_the_measured_window(self): + assert so.tier_verified_for_s(int(208.6 * self.GIB)) == pytest.approx(52.15) + + def test_the_window_grows_with_the_pool_between_its_bounds(self): + sizes = [8, 20, 32, 64, 100, 208.6, 240, 512, 2048] + windows = [so.tier_verified_for_s(int(g * self.GIB)) for g in sizes] + assert windows == sorted(windows) + assert windows[0] == windows[1] == so._TIER_VERIFIED_MIN_S + assert windows[2] == pytest.approx(8.0) + assert windows[-3] == windows[-1] == so._TIER_VERIFIED_MAX_S + + def test_an_unknown_pool_gets_the_floor(self): + assert so.tier_verified_for_s(0) == so._TIER_VERIFIED_MIN_S + + def test_the_index_holds_by_its_own_window(self, clock): + idx = _index(7, pool_bytes=32 * self.GIB, defer_s=2) + assert idx.verified_for_s == pytest.approx(8.0) + clock[0] += 7.5 + assert not idx.needs_verify(7) + clock[0] += 1 + assert idx.needs_verify(7) + + def test_the_block_manager_sizes_the_window_from_the_lmcache_pool( + self, monkeypatch + ): + from test_block_manager import _OFFLOAD_KVC, _bm_with_state_tier + + monkeypatch.setattr( + so, "state_tier_lmcache_sizes", lambda config: (256, 100 * self.GIB) + ) + bm = _bm_with_state_tier(monkeypatch, _OFFLOAD_KVC) + assert bm._joint_chunk_tokens == 256 + assert bm.state_offload.verified_for_s == pytest.approx(25.0) + + # ── the connector: touches out in the metadata, verdicts back by quorum ──── @@ -534,7 +604,7 @@ def _sched(self, connector, idx): def test_voted_hashes_reach_the_connector(self, clock): idx = _index(7) - idx.could_serve(7) + idx.want_touch(7) sent = [] s = self._sched( SimpleNamespace(enqueue_state_touches=lambda hs: sent.extend(hs) or True), @@ -545,11 +615,11 @@ def test_voted_hashes_reach_the_connector(self, clock): def test_a_connector_without_touches_turns_the_refresh_off(self, clock): idx = _index(7) - idx.could_serve(7) + idx.want_touch(7) s = self._sched(SimpleNamespace(), idx) s._publish_state_touches() assert not idx.refresh_hits - clock[0] += so._TIER_VERIFIED_FOR_S + 1 + clock[0] += STALE assert not idx.needs_verify(7), "nothing would ever answer the hold" def test_verdicts_are_applied_to_the_index(self, clock): @@ -609,7 +679,7 @@ def test_an_unconfirmed_joint_image_is_held(self, clock, monkeypatch): bm = _paged_bm() commits = self._joint(bm, monkeypatch) bm.state_offload.note_stored(77) - clock[0] += so._TIER_VERIFIED_FOR_S + 1 + clock[0] += STALE seq = _seq() assert not bm.record_allocation( @@ -630,7 +700,7 @@ def test_a_gone_image_is_not_parked_on(self, clock, monkeypatch): bm = _paged_bm() self._joint(bm, monkeypatch) bm.state_offload.note_stored(77) - clock[0] += so._TIER_VERIFIED_FOR_S + 1 + clock[0] += STALE assert not bm.record_allocation( _seq(), AdmissionProbe(2, [1, 2]), defer_unverified=True ) @@ -651,7 +721,7 @@ def test_a_state_resident_in_hbm_is_never_held(self, clock, monkeypatch): bm = _paged_bm() self._joint(bm, monkeypatch, from_hbm=True) bm.state_offload.note_stored(77) - clock[0] += so._TIER_VERIFIED_FOR_S + 1 + clock[0] += STALE assert bm.record_allocation( _seq(), AdmissionProbe(2, [1, 2]), defer_unverified=True ) @@ -660,7 +730,7 @@ def test_the_fit_probe_path_never_holds(self, clock, monkeypatch): bm = _paged_bm() self._joint(bm, monkeypatch) bm.state_offload.note_stored(77) - clock[0] += so._TIER_VERIFIED_FOR_S + 1 + clock[0] += STALE assert bm.record_allocation(_seq(), AdmissionProbe(2, [1, 2])) def test_an_hbm_hit_holds_only_for_a_tier_sourced_state(self, clock, monkeypatch): @@ -674,7 +744,7 @@ def test_an_hbm_hit_holds_only_for_a_tier_sourced_state(self, clock, monkeypatch ) monkeypatch.setattr(bm, "_commit_joint_boundary", lambda seq, d: None) bm.state_offload.note_stored(2) - clock[0] += so._TIER_VERIFIED_FOR_S + 1 + clock[0] += STALE resident = {2} monkeypatch.setattr( bm.paged_state_checkpoints, "contains", lambda h: h in resident @@ -697,7 +767,7 @@ def test_the_hold_judges_the_joint_fit_kept(self, clock, monkeypatch): ) monkeypatch.setattr(bm.paged_state_checkpoints, "contains", lambda h: h == 2) bm.state_offload.note_stored(77) - clock[0] += so._TIER_VERIFIED_FOR_S + 1 + clock[0] += STALE assert bm.record_allocation( _seq(), AdmissionProbe(2, [1, 2], 1), defer_unverified=True ) @@ -705,6 +775,112 @@ def test_the_hold_judges_the_joint_fit_kept(self, clock, monkeypatch): assert bm.state_offload.loads_deferred_for_verify == 0 +class TestVotersRegisterTouches: + """`could_serve` only answers; the voters register the hits they count.""" + + def _bm(self): + bm = _paged_bm() + bm.paged_state_checkpoints.attach_offload(bm.state_offload) + bm.state_offload.note_stored(77) + return bm + + def test_the_resumable_scan_registers_a_tier_hit(self, clock): + bm = self._bm() + assert bm.paged_state_checkpoints._reachable(77) + assert bm.state_offload.take_touches() == [77] + + def test_an_hbm_hit_is_not_a_tier_vote(self, clock, monkeypatch): + bm = self._bm() + monkeypatch.setattr( + bm.paged_state_checkpoints.store, "contains", lambda h: True + ) + assert bm.paged_state_checkpoints._reachable(77) + assert bm.state_offload.take_touches() == [] + + def test_asking_whether_the_tier_can_serve_registers_nothing(self, clock): + bm = self._bm() + assert bm._tier_can_serve(77) + assert bm.state_offload.take_touches() == [] + + def test_an_offered_load_registers_its_hash(self, clock): + bm = self._bm() + seq = _seq() + assert bm._request_state_load(seq, 77) + assert bm.state_offload.take_touches() == [77] + + +class TestResumeHash: + """`allocate` and the admission hold pick the state to resume from with + one helper, `_state_resume_hash`.""" + + def test_the_joint_boundary_wins_then_a_replay_then_the_hit(self, monkeypatch): + bm = _paged_bm() + seq = _seq() + assert bm._state_resume_hash(seq, 77, 2) == 77 + assert bm._state_resume_hash(seq, -1, 2) == 2 + assert bm._state_resume_hash(seq, -1, -1) == -1 + monkeypatch.setattr(bm, "state_replay_tokens", 8) + seq.replays_hit = True + assert bm._state_resume_hash(seq, -1, 2) == -1, "a replay restores none" + assert bm._state_resume_hash(seq, 77, 2) == 77 + + def test_a_replayed_hit_is_not_held(self, clock, monkeypatch): + from atom.model_engine.block_manager import JointBoundaryDecision + + bm = _paged_bm() + monkeypatch.setattr( + bm, + "_joint_kv_boundary", + lambda *a: JointBoundaryDecision(skip_reason="no_kv_reach"), + ) + monkeypatch.setattr(bm, "_commit_joint_boundary", lambda seq, d: None) + monkeypatch.setattr(bm.paged_state_checkpoints, "contains", lambda h: False) + monkeypatch.setattr(bm, "state_replay_tokens", 8) + bm.state_offload.note_stored(2) + clock[0] += STALE + seq = _seq() + seq.replays_hit = True + assert bm.record_allocation( + seq, AdmissionProbe(2, [1, 2]), defer_unverified=True + ) + + def test_the_hold_judges_the_hash_allocate_then_loads(self, clock, monkeypatch): + from test_state_checkpoint import boundary_hash, stateful_seq + + bm = _paged_bm() + offload = bm.state_offload + bm.paged_state_checkpoints.attach_offload(offload) + first = _seq() + bm.allocate(first, bm.can_allocate(first)) + bm.hash_blocks(first, bm.checkpoint_limit(first) - first.num_cached_tokens) + h = boundary_hash(bm, first) + bm.take_state_maintenance_ops() + bm.complete_previous_state_batch() + # The image left HBM; only the tier has it, and not confirmed lately. + bm.paged_state_checkpoints.unindex(h) + offload.note_stored(h) + clock[0] += STALE + + asked, attached = [], [] + verify, attach = offload.needs_verify, bm._attach_state_slots + monkeypatch.setattr( + offload, "needs_verify", lambda x: asked.append(x) or verify(x) + ) + monkeypatch.setattr( + bm, "_attach_state_slots", lambda s, x: attached.append(x) or attach(s, x) + ) + second = stateful_seq(list(range(48))) + probe = AdmissionProbe() + probe.num_cached_blocks = bm.can_allocate(second, record=False, probe=probe) + assert probe.num_cached_blocks > 0 + assert not bm.record_allocation(second, probe, defer_unverified=True) + offload.apply_touch_verdict(h, True) + assert bm.record_allocation(second, probe, defer_unverified=True) + assert bm.allocate(second, probe.num_cached_blocks) + assert set(asked) == {h} and attached == [h] + assert second.offload_joint.load_hash == h + + def test_the_scheduler_retries_a_held_admission(monkeypatch): """A held seq stays waiting and is admitted on a later pass.""" from test_state_checkpoint import ckpt_config, make_scheduler, stateful_seq From 9219a8dba216aab253ef0be2dab86b3f2ad36781 Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Fri, 30 Oct 2026 14:11:14 +0000 Subject: [PATCH 06/18] fix(offload): count and warn on presence-only touches, tolerate an offload without want_touch A touch that finds the image outside LocalCPU's hot cache only confirms presence; count it (unrefreshed_touches) and warn once. The verdict stays resident: get can read the image now, and a later miss degrades to a recompute. PagedStateCheckpointCoordinator._reachable registers the touch through getattr, like stats, so an attached offload without want_touch keeps its tier hits. --- .../offload/hybrid/kimi_k3/state_object.py | 18 ++++++++- atom/model_engine/page_unit_checkpoint.py | 7 +++- tests/test_state_tier_lru_refresh.py | 37 +++++++++++++++++++ 3 files changed, 59 insertions(+), 3 deletions(-) diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py index 122859c7d9..4a423b1fa2 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py @@ -57,6 +57,7 @@ def __init__( self._worker_id = int(worker_id) self._layout_id = layout_id self._misfit_reads = 0 + self._unrefreshed_touches = 0 self._storage = None self._max_entries = ( None @@ -190,8 +191,21 @@ def _refresh(self, key) -> bool: policy.update_on_hit(key, hot) return True # Not in the hot cache (or there is none): another backend may still - # hold it, and `get` reads every backend. Presence only. - return bool(self._storage.contains(key)) + # hold it, and `get` reads every backend. Presence only. Still reported + # resident: `get` can read it now, and the confirmation that opens only + # lets later loads skip the hold, where a miss degrades to a recompute. + if not self._storage.contains(key): + return False + self._unrefreshed_touches += 1 + if self._unrefreshed_touches == 1: + logger.warning( + "state offload: a touched state image is outside LocalCPU's " + "hot cache (or no LocalCPU LRU is reachable), so the touch " + "only confirmed it is present and did not refresh its " + "eviction order (counted as unrefreshed_touches=%d)", + self._unrefreshed_touches, + ) + return True def _put(self, h: int, unit_ids, on_source_released=None) -> bool: """Store one checkpoint image. False when nothing was stored. diff --git a/atom/model_engine/page_unit_checkpoint.py b/atom/model_engine/page_unit_checkpoint.py index 2adb966c81..00051ffedd 100644 --- a/atom/model_engine/page_unit_checkpoint.py +++ b/atom/model_engine/page_unit_checkpoint.py @@ -1080,7 +1080,12 @@ def _reachable(self, h: int) -> bool: return True if self.offload is None or not self.offload.could_serve(h): return False - self.offload.want_touch(h) + # `getattr`, as for `stats`: `attach_offload` accepts anything that + # answers `could_serve`, and one without touches loses the LRU refresh, + # not the hit. + want_touch = getattr(self.offload, "want_touch", None) + if callable(want_touch): + want_touch(h) return True def checkpoint(self, seq: Sequence, boundary_blocks: int, h: int) -> None: diff --git a/tests/test_state_tier_lru_refresh.py b/tests/test_state_tier_lru_refresh.py index 9c915fb4fa..8e045b25f4 100644 --- a/tests/test_state_tier_lru_refresh.py +++ b/tests/test_state_tier_lru_refresh.py @@ -16,6 +16,7 @@ """ import dataclasses +import logging import sys import threading import types @@ -281,6 +282,33 @@ def test_without_a_local_cpu_backend_a_touch_asks_presence(self, lmcache_key): pool.hot.clear() assert codec.touch(1) is False + def test_a_presence_only_touch_is_counted_and_warned_once( + self, lmcache_key, caplog + ): + """It confirms the image without refreshing it, so it must not be + silent.""" + pool = _LocalCPU() + codec, _ = _codec(pool, max_entries=None) + assert codec.put(1, []) + pool.local_cpu_backend = None + with caplog.at_level(logging.WARNING, logger="atom"): + assert codec.touch(1) is True + assert codec.touch(1) is True + assert codec._unrefreshed_touches == 2 + warned = [r for r in caplog.records if "unrefreshed_touches" in r.message] + assert len(warned) == 1 and "unrefreshed_touches=1" in warned[0].message + + def test_a_refreshed_or_gone_image_is_not_counted(self, lmcache_key, caplog): + pool = _LocalCPU() + codec, _ = _codec(pool, max_entries=None) + assert codec.put(1, []) + with caplog.at_level(logging.WARNING, logger="atom"): + assert codec.touch(1) is True + pool.hot.clear() + assert codec.touch(1) is False + assert codec._unrefreshed_touches == 0 + assert not [r for r in caplog.records if "unrefreshed_touches" in r.message] + # ── the index: votes queue touches, verdicts keep it in step ─────────────── @@ -808,6 +836,15 @@ def test_an_offered_load_registers_its_hash(self, clock): assert bm._request_state_load(seq, 77) assert bm.state_offload.take_touches() == [77] + def test_an_offload_without_touches_still_counts_a_tier_hit(self, clock): + """`attach_offload` takes any index answering `could_serve`.""" + bm = _paged_bm() + bm.paged_state_checkpoints.attach_offload( + SimpleNamespace(hashes={77}, could_serve=lambda h: h == 77) + ) + assert bm.paged_state_checkpoints._reachable(77) + assert not bm.paged_state_checkpoints._reachable(78) + class TestResumeHash: """`allocate` and the admission hold pick the state to resume from with From 1b311b688f9b409cf66826c8b2e6b6832f33e89f Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Fri, 30 Oct 2026 15:36:10 +0000 Subject: [PATCH 07/18] state offload: resolve LMCache's LRU at bind, expose unrefreshed touches - bind_storage_manager resolves LocalCPUBackend's cpu_lock / hot_cache / cache_policy.update_on_hit once and raises when one is missing, instead of a per-touch getattr that silently degraded to contains(). A storage manager with no LocalCPU backend warns once and every touch is counted. - put()'s skip-the-store fast path only applies when the image is in LocalCPU's hot cache; a disk-only image is stored again so it returns to the CPU tier. - touch() returns TouchResult (GONE / PRESENT / REFRESHED). The worker reports refreshed-ness on a new k3_state_touch_refresh channel (failure-dominant across TP) and the scheduler surfaces it as state_offload_touches_unrefreshed in the state checkpoints stats line. --- .../disaggregation/multi/multi_connector.py | 7 + atom/kv_transfer/offload/connector.py | 5 + .../offload/hybrid/kimi_k3/connector.py | 43 +++- .../offload/hybrid/kimi_k3/state_object.py | 132 +++++++---- atom/model_engine/scheduler.py | 3 + atom/model_engine/state_offload.py | 8 +- tests/test_lmcache_offload_connector.py | 14 +- tests/test_state_tier_lru_refresh.py | 217 +++++++++++++++--- 8 files changed, 352 insertions(+), 77 deletions(-) diff --git a/atom/kv_transfer/disaggregation/multi/multi_connector.py b/atom/kv_transfer/disaggregation/multi/multi_connector.py index f7c7ff3b0c..a7cefe9ee5 100644 --- a/atom/kv_transfer/disaggregation/multi/multi_connector.py +++ b/atom/kv_transfer/disaggregation/multi/multi_connector.py @@ -678,6 +678,13 @@ def take_state_touch_verdicts(self) -> dict: 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. diff --git a/atom/kv_transfer/offload/connector.py b/atom/kv_transfer/offload/connector.py index 9c3c4052cb..352c0fa697 100644 --- a/atom/kv_transfer/offload/connector.py +++ b/atom/kv_transfer/offload/connector.py @@ -348,6 +348,11 @@ def take_state_touch_verdicts(self) -> dict: 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 diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py index ab527b4cf2..751d8b785d 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py @@ -31,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, @@ -70,6 +73,11 @@ #: across TP, like a joint load: one missing shard means the load would miss. STATE_TOUCH_CHANNEL = "k3_state_touch" +#: Alongside each `STATE_TOUCH_CHANNEL` verdict, same id: `succeeded=False` when +#: the image was present but outside LocalCPU's hot cache, so the touch did not +#: refresh its eviction order. Failure-dominant: any such shard counts. +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, @@ -124,8 +132,8 @@ 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`. - self._touch_verdicts: dict[StateTouchId, bool] = {} + # Touch verdicts since the last `get_finished`: (resident, refreshed). + self._touch_verdicts: dict[StateTouchId, tuple[bool, bool]] = {} def close(self) -> None: """Drain then join the state tier before the base executors. @@ -318,10 +326,10 @@ def _touch_states(self, metadata) -> None: tier = self._state_tier for touch_id in touches: if tier is None: - self._touch_verdicts[touch_id] = False + self._touch_verdicts[touch_id] = (False, True) continue try: - resident = bool(tier.codec.touch(int(touch_id.prefix_hash))) + result = tier.codec.touch(int(touch_id.prefix_hash)) except Exception: # Says nothing about the bytes; do not let it erase the index. logger.warning( @@ -329,8 +337,11 @@ def _touch_states(self, metadata) -> None: touch_id.prefix_hash, exc_info=True, ) - resident = True - self._touch_verdicts[touch_id] = resident + result = TouchResult.REFRESHED + 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. @@ -513,10 +524,13 @@ def get_finished(self): 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 in self._touch_verdicts.items(): + 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 @@ -666,6 +680,8 @@ def __init__(self, config) -> None: self._pending_state_touches: list[StateTouchId] = [] self._touch_round = 0 self._state_touch_verdicts: dict[int, bool] = {} + # Touches that only confirmed presence (`take_state_touches_unrefreshed`). + 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 @@ -722,6 +738,13 @@ def take_state_touch_verdicts(self) -> dict[int, bool]: self._state_touch_verdicts = {} return verdicts + def take_state_touches_unrefreshed(self) -> int: + """Touches since the last call that found an image present but could + not refresh it in LocalCPU's LRU.""" + n = self._state_touches_unrefreshed + self._state_touches_unrefreshed = 0 + return n + def build_connector_meta(self) -> LMCacheOffloadMetadata: self._refresh_save_stall() meta = super().build_connector_meta() @@ -948,6 +971,10 @@ def connector_completion(self, completion) -> bool | None: # 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 diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py index 4a423b1fa2..f000235b38 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py @@ -11,6 +11,7 @@ from __future__ import annotations +import enum import logging import threading from collections import OrderedDict @@ -21,6 +22,56 @@ logger = logging.getLogger("atom") +class TouchResult(enum.IntEnum): + """What a touch found. Truthy when the image can still be read.""" + + GONE = 0 + #: Outside LocalCPU's hot cache but held by another backend (`get` reads + #: it and promotes it back): present, its eviction order not refreshed. + PRESENT = 1 + REFRESHED = 2 + + +def _local_cpu_lru(storage_manager): + """LMCache's LocalCPU backend, checked once for what a touch needs, or + None when the storage manager has none. + + `cpu_lock`, `hot_cache` and `cache_policy.update_on_hit` are LMCache + internals (written against 0.5.6), not API. A LocalCPU backend missing one + is an LMCache this code does not know, so raise: degrading every touch to + a presence check would keep every load "succeeding" while no image is + refreshed. + """ + if not hasattr(storage_manager, "local_cpu_backend"): + raise RuntimeError( + "state offload: LMCache's storage manager has no `local_cpu_backend`; " + "state images cannot be refreshed in its LRU" + ) + backend = storage_manager.local_cpu_backend + if backend is None: + logger.warning( + "state offload: LMCache has no LocalCPU backend, so state touches " + "only confirm presence and refresh no eviction order (counted as " + "state_offload_touches_unrefreshed)" + ) + return None + missing = [ + name + for name in ("cpu_lock", "hot_cache", "cache_policy") + if getattr(backend, name, None) is None + ] + if not missing and not callable( + getattr(backend.cache_policy, "update_on_hit", None) + ): + missing.append("cache_policy.update_on_hit") + if missing: + raise RuntimeError( + "state offload: LMCache's LocalCPUBackend has no " + f"{', '.join(missing)}; state images cannot be refreshed in its LRU" + ) + return backend + + class StateByteCodec: """Pack/unpack one entry's state and move it through `storage_manager`. @@ -59,6 +110,7 @@ def __init__( self._misfit_reads = 0 self._unrefreshed_touches = 0 self._storage = None + self._local_cpu = None self._max_entries = ( None if max_cache_bytes is None @@ -83,6 +135,7 @@ def __init__( ) def bind_storage_manager(self, storage_manager) -> None: + self._local_cpu = _local_cpu_lru(storage_manager) self._storage = storage_manager def key(self, h: int): @@ -127,10 +180,12 @@ def put(self, h: int, unit_ids, on_source_released=None) -> bool: if self._storage is None: return False key = self.key(h) - if self._touch_key(key): - # Same prefix hash and layout, same bytes, still resident: storing - # again would only cost an image-sized D2H that LocalCPU then - # refuses as a duplicate key. The touch counts as the use it is. + if self._touch_key(key) is TouchResult.REFRESHED: + # Same prefix hash and layout, same bytes, still in LocalCPU's hot + # cache: storing again would only cost an image-sized D2H that + # LocalCPU then refuses as a duplicate key. The touch counts as the + # use it is. Only there: an image another backend holds (disk) is + # stored again, which puts it back into the CPU tier. if on_source_released is not None: on_source_released() return True @@ -150,8 +205,8 @@ def put(self, h: int, unit_ids, on_source_released=None) -> bool: self._cache_keys[key] = None return ok - def touch(self, h: int) -> bool: - """Refresh `h`'s image in LMCache's LRU; False when it is gone. + def touch(self, h: int) -> TouchResult: + """Refresh `h`'s image in LMCache's LRU; falsy (`GONE`) when it is gone. The state-image twin of what a KV lookup does to its chunks (`LocalCPUBackend.batched_async_contains(pin=True)` then @@ -166,46 +221,45 @@ def touch(self, h: int) -> bool: is the list the KV lookup is filling at the same time. """ if self._storage is None: - return False - return self._touch_key(self.key(h)) - - def _touch_key(self, key) -> bool: - resident = self._refresh(key) + return TouchResult.GONE + result = self._touch_key(self.key(h)) + if result is TouchResult.PRESENT: + # Still reported resident: `get` can read it now, and the + # confirmation that opens only lets later loads skip the hold, + # where a miss degrades to a recompute. + self._unrefreshed_touches += 1 + if self._unrefreshed_touches == 1: + logger.warning( + "state offload: a touched state image is outside " + "LocalCPU's hot cache, so the touch only confirmed it is " + "present and did not refresh its eviction order (counted " + "as state_offload_touches_unrefreshed)" + ) + return result + + def _touch_key(self, key) -> TouchResult: + result = self._refresh(key) with self._cache_lock: - if resident: + if result: if key in self._cache_keys: self._cache_keys.move_to_end(key) else: # Evicted by LMCache: it no longer counts against the cap. self._cache_keys.pop(key, None) - return resident - - def _refresh(self, key) -> bool: - backend = getattr(self._storage, "local_cpu_backend", None) - lock = getattr(backend, "cpu_lock", None) - hot = getattr(backend, "hot_cache", None) - policy = getattr(backend, "cache_policy", None) - if lock is not None and hot is not None and policy is not None: - with lock: - if key in hot: - policy.update_on_hit(key, hot) - return True + return result + + def _refresh(self, key) -> TouchResult: + backend = self._local_cpu + if backend is not None: + with backend.cpu_lock: + if key in backend.hot_cache: + backend.cache_policy.update_on_hit(key, backend.hot_cache) + return TouchResult.REFRESHED # Not in the hot cache (or there is none): another backend may still - # hold it, and `get` reads every backend. Presence only. Still reported - # resident: `get` can read it now, and the confirmation that opens only - # lets later loads skip the hold, where a miss degrades to a recompute. - if not self._storage.contains(key): - return False - self._unrefreshed_touches += 1 - if self._unrefreshed_touches == 1: - logger.warning( - "state offload: a touched state image is outside LocalCPU's " - "hot cache (or no LocalCPU LRU is reachable), so the touch " - "only confirmed it is present and did not refresh its " - "eviction order (counted as unrefreshed_touches=%d)", - self._unrefreshed_touches, - ) - return True + # hold it, and `get` reads every backend. Presence only. + if self._storage.contains(key): + return TouchResult.PRESENT + return TouchResult.GONE def _put(self, h: int, unit_ids, on_source_released=None) -> bool: """Store one checkpoint image. False when nothing was stored. diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index 68c370e0ae..e49b5b8e48 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -3811,6 +3811,9 @@ def _apply_state_touch_verdicts(self, offload) -> None: return for h, resident in take().items(): offload.apply_touch_verdict(h, resident) + unrefreshed = getattr(self.kv_connector, "take_state_touches_unrefreshed", None) + if unrefreshed is not None: + offload.touches_unrefreshed += int(unrefreshed()) def _publish_state_loads(self) -> None: """Hand this pass's state-tier loads to the connector, before it builds diff --git a/atom/model_engine/state_offload.py b/atom/model_engine/state_offload.py index 36631ff488..325701f53c 100644 --- a/atom/model_engine/state_offload.py +++ b/atom/model_engine/state_offload.py @@ -138,6 +138,9 @@ def __init__( # Touches that found the image gone: LMCache evicted it and the index # learned of it before any request parked on it. self.touch_misses = 0 + # Touches that found the image present but outside LocalCPU's hot + # cache, so they confirmed it without refreshing its eviction order. + self.touches_unrefreshed = 0 # Admissions held back a pass because the image they would load had # not been confirmed recently (`needs_verify`). self.loads_deferred_for_verify = 0 @@ -403,9 +406,12 @@ def stats(self) -> dict[str, int]: "hashes_evicted": self.hashes_evicted, # Hit refresh. `touch_misses` are images LMCache dropped that left # the index before anything parked on them; `loads_failed` is what - # still got past. + # still got past. `touches_unrefreshed` is the part of the + # resident verdicts that only confirmed presence: 0 when every + # touch refreshed LMCache's LRU. "touches_sent": self.touches_sent, "touch_misses": self.touch_misses, + "touches_unrefreshed": self.touches_unrefreshed, "loads_deferred_for_verify": self.loads_deferred_for_verify, "verify_holds_expired": self.verify_holds_expired, } diff --git a/tests/test_lmcache_offload_connector.py b/tests/test_lmcache_offload_connector.py index 9c68522edc..a6546babcf 100644 --- a/tests/test_lmcache_offload_connector.py +++ b/tests/test_lmcache_offload_connector.py @@ -7,7 +7,7 @@ import sys import threading import types -from collections import deque +from collections import OrderedDict, deque from contextlib import nullcontext from types import SimpleNamespace @@ -6160,10 +6160,18 @@ def test_no_more_saves_go_out_than_the_pool_can_afford_to_pin(): class _FakeStorage: - """The two `storage_manager` calls the codec makes, over a plain dict.""" + """The `storage_manager` calls the codec makes, over a plain dict, with + the LocalCPU LRU internals `bind_storage_manager` requires.""" def __init__(self) -> None: - self.objects: dict = {} + self.objects: OrderedDict = OrderedDict() + self.local_cpu_backend = SimpleNamespace( + cpu_lock=threading.Lock(), + hot_cache=self.objects, + cache_policy=SimpleNamespace( + update_on_hit=lambda key, cache: cache.move_to_end(key) + ), + ) def batched_put(self, keys, objs) -> None: for k, o in zip(keys, objs, strict=True): diff --git a/tests/test_state_tier_lru_refresh.py b/tests/test_state_tier_lru_refresh.py index 8e045b25f4..b1fc2a66e2 100644 --- a/tests/test_state_tier_lru_refresh.py +++ b/tests/test_state_tier_lru_refresh.py @@ -26,6 +26,7 @@ import pytest from atom.kv_transfer.disaggregation.types import ConnectorCompletion, StateTouchId +from atom.kv_transfer.offload.hybrid.kimi_k3.state_object import TouchResult from atom.model_engine import state_offload as so from atom.model_engine.block_manager import AdmissionProbe from atom.model_engine.state_offload import ( @@ -102,8 +103,12 @@ class _LocalCPU: def __init__(self, refuse_all: bool = False) -> None: self.hot: OrderedDict = OrderedDict() + # A second backend (LocalDisk): read by `get`, which promotes a hit + # back into the hot cache, and answered by `contains`. + self.disk: dict = {} self.refuse_all = refuse_all self.fail_put = False + self.puts = 0 self.local_cpu_backend = SimpleNamespace( cpu_lock=threading.Lock(), hot_cache=self.hot, cache_policy=_LRUPolicy() ) @@ -111,6 +116,7 @@ def __init__(self, refuse_all: bool = False) -> None: def batched_put(self, keys, objs) -> None: if self.fail_put: raise RuntimeError("put failed") + self.puts += 1 for k, o in zip(keys, objs, strict=True): if not self.refuse_all and k not in self.hot: o.ref_count_up() @@ -124,7 +130,14 @@ def get(self, key): return o def contains(self, key): - return "LocalCPUBackend" if key in self.hot else None + if key in self.hot: + return "LocalCPUBackend" + return "LocalDiskBackend" if key in self.disk else None + + def spill_to_disk(self, key) -> None: + """LocalCPU evicted `key`; the disk backend still holds a copy.""" + self.disk[key] = True + self.hot.pop(key).ref_count_down() def remove(self, key) -> int: o = self.hot.pop(key, None) @@ -189,7 +202,7 @@ def test_a_touch_moves_the_image_past_older_kv(self, lmcache_key): codec, _ = _codec(pool, max_entries=None) assert codec.put(1, []) _kv(pool, "kv-chunk") - assert codec.touch(1) is True + assert codec.touch(1) is TouchResult.REFRESHED assert pool.evict_oldest() == "kv-chunk" assert codec.key(1) in pool.hot @@ -198,7 +211,7 @@ def test_a_touch_reports_an_evicted_image_gone(self, lmcache_key): codec, _ = _codec(pool, max_entries=2) assert codec.put(1, []) assert pool.evict_oldest() == codec.key(1) - assert codec.touch(1) is False + assert codec.touch(1) is TouchResult.GONE assert codec.key(1) not in codec._cache_keys, "no longer counts against the cap" def test_a_resume_refreshes_the_image(self, lmcache_key): @@ -244,7 +257,7 @@ def test_the_cap_retires_the_least_recently_used_image(self, lmcache_key): pool = _LocalCPU() codec, objs = _codec(pool, max_entries=2) assert codec.put(1, []) and codec.put(2, []) - assert codec.touch(1) + assert codec.touch(1) is TouchResult.REFRESHED assert codec.put(3, []) assert codec.key(2) not in pool.hot @@ -273,41 +286,91 @@ def test_a_failed_put_discharges_its_reference(self, lmcache_key): codec.put(1, []) assert objs[0].refs == 0 - def test_without_a_local_cpu_backend_a_touch_asks_presence(self, lmcache_key): + def test_without_a_local_cpu_backend_a_touch_asks_presence( + self, lmcache_key, caplog + ): pool = _LocalCPU() - codec, _ = _codec(pool, max_entries=None) - assert codec.put(1, []) pool.local_cpu_backend = None - assert codec.touch(1) is True + with caplog.at_level(logging.WARNING, logger="atom"): + codec, _ = _codec(pool, max_entries=None) + assert any("no LocalCPU backend" in r.message for r in caplog.records) + pool.hot[codec.key(1)] = _Obj() + assert codec.touch(1) is TouchResult.PRESENT pool.hot.clear() - assert codec.touch(1) is False + assert codec.touch(1) is TouchResult.GONE - def test_a_presence_only_touch_is_counted_and_warned_once( - self, lmcache_key, caplog + @pytest.mark.parametrize("dropped", ["cpu_lock", "hot_cache", "cache_policy"]) + def test_a_local_cpu_without_the_lru_internals_fails_at_bind( + self, lmcache_key, dropped ): - """It confirms the image without refreshing it, so it must not be - silent.""" + """An LMCache whose names moved must not degrade every touch to a + presence check that keeps every load "succeeding".""" + pool = _LocalCPU() + setattr(pool.local_cpu_backend, dropped, None) + with pytest.raises(RuntimeError, match=dropped): + _codec(pool, max_entries=None) + + def test_a_policy_without_update_on_hit_fails_at_bind(self, lmcache_key): + pool = _LocalCPU() + pool.local_cpu_backend.cache_policy = SimpleNamespace() + with pytest.raises(RuntimeError, match="update_on_hit"): + _codec(pool, max_entries=None) + + def test_a_storage_manager_without_local_cpu_backend_fails_at_bind( + self, lmcache_key + ): + pool = _LocalCPU() + del pool.local_cpu_backend + with pytest.raises(RuntimeError, match="local_cpu_backend"): + _codec(pool, max_entries=None) + + def test_the_lru_is_resolved_once_at_bind(self, lmcache_key): pool = _LocalCPU() codec, _ = _codec(pool, max_entries=None) assert codec.put(1, []) + _kv(pool, "kv-chunk") pool.local_cpu_backend = None + assert codec.touch(1) is TouchResult.REFRESHED + assert pool.evict_oldest() == "kv-chunk" + + def test_a_disk_only_image_is_counted_and_warned_once(self, lmcache_key, caplog): + """Present, so still readable, but its LRU was not refreshed: it must + not be silent.""" + pool = _LocalCPU() + codec, _ = _codec(pool, max_entries=None) + assert codec.put(1, []) + pool.spill_to_disk(codec.key(1)) with caplog.at_level(logging.WARNING, logger="atom"): - assert codec.touch(1) is True - assert codec.touch(1) is True + assert codec.touch(1) is TouchResult.PRESENT + assert codec.touch(1) is TouchResult.PRESENT assert codec._unrefreshed_touches == 2 - warned = [r for r in caplog.records if "unrefreshed_touches" in r.message] - assert len(warned) == 1 and "unrefreshed_touches=1" in warned[0].message + warned = [r for r in caplog.records if "touches_unrefreshed" in r.message] + assert len(warned) == 1 def test_a_refreshed_or_gone_image_is_not_counted(self, lmcache_key, caplog): pool = _LocalCPU() codec, _ = _codec(pool, max_entries=None) assert codec.put(1, []) with caplog.at_level(logging.WARNING, logger="atom"): - assert codec.touch(1) is True + assert codec.touch(1) is TouchResult.REFRESHED pool.hot.clear() - assert codec.touch(1) is False + assert codec.touch(1) is TouchResult.GONE assert codec._unrefreshed_touches == 0 - assert not [r for r in caplog.records if "unrefreshed_touches" in r.message] + assert not [r for r in caplog.records if "touches_unrefreshed" in r.message] + + def test_a_disk_only_image_is_stored_back_into_the_cpu_tier(self, lmcache_key): + """The put fast path is for an image LocalCPU still holds, not for + one only the disk does.""" + pool = _LocalCPU() + codec, objs = _codec(pool, max_entries=None) + assert codec.put(1, []) + pool.spill_to_disk(codec.key(1)) + released = [] + assert codec.put(1, [], on_source_released=lambda: released.append(True)) + assert pool.puts == 2 and len(objs) == 2 + assert codec.key(1) in pool.hot + assert released == [True] + assert codec._unrefreshed_touches == 0, "a store is not a touch verdict" # ── the index: votes queue touches, verdicts keep it in step ─────────────── @@ -544,20 +607,28 @@ class TestConnectorTouches: def test_the_worker_touches_and_records_each_verdict(self): touched = [] + found = {1: TouchResult.REFRESHED, 2: TouchResult.GONE, 3: TouchResult.PRESENT} + def touch(h): touched.append(h) - return h != 2 + return found[h] worker = _k3_worker(SimpleNamespace(codec=SimpleNamespace(touch=touch))) - ids = [StateTouchId(1, 0), StateTouchId(2, 1)] + ids = [StateTouchId(1, 0), StateTouchId(2, 1), StateTouchId(3, 2)] worker._touch_states(SimpleNamespace(state_touches=ids)) - assert touched == [1, 2] - assert worker._touch_verdicts == {ids[0]: True, ids[1]: False} + assert touched == [1, 2, 3] + # (resident, refreshed): only a present-but-unrefreshed image is False + # on the second. + assert worker._touch_verdicts == { + ids[0]: (True, True), + ids[1]: (False, True), + ids[2]: (True, False), + } def test_with_no_tier_every_image_is_gone(self): worker = _k3_worker(None) worker._touch_states(SimpleNamespace(state_touches=[StateTouchId(1, 0)])) - assert worker._touch_verdicts == {StateTouchId(1, 0): False} + assert worker._touch_verdicts == {StateTouchId(1, 0): (False, True)} def test_a_failing_touch_does_not_erase_the_index(self): def touch(h): @@ -565,7 +636,7 @@ def touch(h): worker = _k3_worker(SimpleNamespace(codec=SimpleNamespace(touch=touch))) worker._touch_states(SimpleNamespace(state_touches=[StateTouchId(1, 0)])) - assert worker._touch_verdicts == {StateTouchId(1, 0): True} + assert worker._touch_verdicts == {StateTouchId(1, 0): (True, True)} def _scheduler(self, monkeypatch): from atom.kv_transfer.offload.dense.connector import DenseOffloadScheduler @@ -585,6 +656,7 @@ def _scheduler(self, monkeypatch): s._pending_state_touches = [] s._touch_round = 0 s._state_touch_verdicts = {} + s._state_touches_unrefreshed = 0 s._refresh_save_stall = lambda: None return s @@ -613,6 +685,85 @@ def test_one_rank_missing_the_image_makes_it_gone(self, monkeypatch): assert s.take_state_touch_verdicts() == {5: False, 6: True} assert s.take_state_touch_verdicts() == {} + def test_a_presence_only_shard_counts_the_touch_unrefreshed(self, monkeypatch): + from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( + STATE_TOUCH_REFRESH_CHANNEL, + ) + + s = self._scheduler(monkeypatch) + for op, ok in ((StateTouchId(5, 0), True), (StateTouchId(6, 1), False)): + handled = s.connector_completion( + ConnectorCompletion(STATE_TOUCH_REFRESH_CHANNEL, op, ok) + ) + assert handled is None, "handled, but not a terminal save" + assert s.take_state_touches_unrefreshed() == 1 + assert s.take_state_touches_unrefreshed() == 0 + assert s.take_state_touch_verdicts() == {}, "not a residency verdict" + + def test_both_touch_reports_reach_quorum_across_tp(self, monkeypatch): + """Each rank files a residency and a refresh report per touch; one + rank that only found the image present counts the touch unrefreshed + without making it gone.""" + from atom.kv_transfer.disaggregation.aggregator import KVOutputAggregator + from atom.kv_transfer.disaggregation.types import KVConnectorOutput + from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( + STATE_TOUCH_CHANNEL, + STATE_TOUCH_REFRESH_CHANNEL, + ) + + touch = StateTouchId(5, 0) + outputs = [] + for result in (TouchResult.REFRESHED, TouchResult.PRESENT): + worker = _k3_worker( + SimpleNamespace(codec=SimpleNamespace(touch=lambda h, r=result: r)) + ) + worker._touch_states(SimpleNamespace(state_touches=[touch])) + resident, refreshed = worker._touch_verdicts[touch] + outputs.append( + KVConnectorOutput( + connector_completions={ + ConnectorCompletion(STATE_TOUCH_CHANNEL, touch, resident), + ConnectorCompletion( + STATE_TOUCH_REFRESH_CHANNEL, touch, refreshed + ), + } + ) + ) + merged = KVOutputAggregator(world_size=2).aggregate(outputs) + s = self._scheduler(monkeypatch) + for completion in merged.connector_completions: + s.connector_completion(completion) + assert s.take_state_touch_verdicts() == {5: True} + assert s.take_state_touches_unrefreshed() == 1 + + def test_the_shells_forward_the_unrefreshed_count(self, monkeypatch): + from atom.kv_transfer.disaggregation.multi.multi_connector import ( + MultiConnectorScheduler, + ) + from atom.kv_transfer.offload.connector import ( + LMCacheOffloadConnectorScheduler, + ) + + shell = LMCacheOffloadConnectorScheduler.__new__( + LMCacheOffloadConnectorScheduler + ) + shell._impl = SimpleNamespace(take_state_touches_unrefreshed=lambda: 2) + assert shell.take_state_touches_unrefreshed() == 2 + shell._impl = SimpleNamespace() + assert shell.take_state_touches_unrefreshed() == 0 + + multi = MultiConnectorScheduler.__new__(MultiConnectorScheduler) + monkeypatch.setattr( + MultiConnectorScheduler, + "_state_tier_sub", + lambda self: SimpleNamespace(take_state_touches_unrefreshed=lambda: 4), + ) + assert multi.take_state_touches_unrefreshed() == 4 + monkeypatch.setattr( + MultiConnectorScheduler, "_state_tier_sub", lambda self: None + ) + assert multi.take_state_touches_unrefreshed() == 0 + def test_a_negative_round_is_refused(self): with pytest.raises(ValueError): StateTouchId(1, -1) @@ -659,6 +810,20 @@ def test_verdicts_are_applied_to_the_index(self, clock): s._apply_state_touch_verdicts(idx) assert not idx.could_serve(7) and idx.could_serve(8) + def test_presence_only_touches_reach_the_stats_line(self, clock): + idx = _index(7) + s = self._sched( + SimpleNamespace( + take_state_touch_verdicts=lambda: {7: True}, + take_state_touches_unrefreshed=lambda: 3, + ), + idx, + ) + assert idx.stats()["touches_unrefreshed"] == 0 + s._apply_state_touch_verdicts(idx) + s._apply_state_touch_verdicts(idx) + assert idx.stats()["touches_unrefreshed"] == 6 + def _paged_bm(): from test_state_checkpoint import ( From 4eff58198be6788bff61b733d8c65e33709dde00 Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Fri, 30 Oct 2026 16:18:03 +0000 Subject: [PATCH 08/18] state offload: degrade on missing LRU internals, drop stale gone verdicts, no hold on a store-only tier - A storage manager missing local_cpu_backend or one of its LRU internals no longer raises at bind: touches only confirm presence, bind warns once naming what is missing, and every touch is counted in state_offload_touches_unrefreshed. - apply_touch_verdict ignores a gone verdict when the hash was confirmed (stored) after its latest touch was sent; the touch looked before those bytes landed. - needs_verify asks could_serve, so a store-only role (kv_producer) does not hold an admission for a load it never makes. --- atom/kv_transfer/offload/README.md | 10 ++- .../offload/hybrid/kimi_k3/state_object.py | 58 ++++++------- atom/model_engine/state_offload.py | 9 +- tests/test_state_tier_lru_refresh.py | 87 ++++++++++++++++--- 4 files changed, 115 insertions(+), 49 deletions(-) diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index 4d7e759e7c..d687674fc2 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -129,8 +129,13 @@ every hit the index votes for is also sent to the workers as a *touch* 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 admission about to load an image not confirmed recently is -held for a pass (at most `OFFLOAD_LOOKUP_DEFER_S`) while that verdict comes +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 @@ -1461,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) diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py index f000235b38..4ae1dc39ae 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py @@ -34,42 +34,36 @@ class TouchResult(enum.IntEnum): def _local_cpu_lru(storage_manager): """LMCache's LocalCPU backend, checked once for what a touch needs, or - None when the storage manager has none. + None when a touch cannot refresh anything. `cpu_lock`, `hot_cache` and `cache_policy.update_on_hit` are LMCache - internals (written against 0.5.6), not API. A LocalCPU backend missing one - is an LMCache this code does not know, so raise: degrading every touch to - a presence check would keep every load "succeeding" while no image is - refreshed. + internals (written against 0.5.6), not API. Without one of them, touches + only confirm presence, like a storage manager with no LocalCPU backend: + loads still work, only the refresh is lost. Not silent: it warns here and + every touch is counted in `state_offload_touches_unrefreshed`. """ - if not hasattr(storage_manager, "local_cpu_backend"): - raise RuntimeError( - "state offload: LMCache's storage manager has no `local_cpu_backend`; " - "state images cannot be refreshed in its LRU" - ) - backend = storage_manager.local_cpu_backend + backend = getattr(storage_manager, "local_cpu_backend", None) if backend is None: - logger.warning( - "state offload: LMCache has no LocalCPU backend, so state touches " - "only confirm presence and refresh no eviction order (counted as " - "state_offload_touches_unrefreshed)" - ) - return None - missing = [ - name - for name in ("cpu_lock", "hot_cache", "cache_policy") - if getattr(backend, name, None) is None - ] - if not missing and not callable( - getattr(backend.cache_policy, "update_on_hit", None) - ): - missing.append("cache_policy.update_on_hit") - if missing: - raise RuntimeError( - "state offload: LMCache's LocalCPUBackend has no " - f"{', '.join(missing)}; state images cannot be refreshed in its LRU" - ) - return backend + missing = ["local_cpu_backend"] + else: + missing = [ + f"local_cpu_backend.{name}" + for name in ("cpu_lock", "hot_cache", "cache_policy") + if getattr(backend, name, None) is None + ] + if not missing and not callable( + getattr(backend.cache_policy, "update_on_hit", None) + ): + missing.append("local_cpu_backend.cache_policy.update_on_hit") + if not missing: + return backend + logger.warning( + "state offload: LMCache's storage manager has no %s, so state touches " + "only confirm presence and refresh no eviction order (every touch is " + "counted in state_offload_touches_unrefreshed)", + ", ".join(missing), + ) + return None class StateByteCodec: diff --git a/atom/model_engine/state_offload.py b/atom/model_engine/state_offload.py index 325701f53c..92ba77638a 100644 --- a/atom/model_engine/state_offload.py +++ b/atom/model_engine/state_offload.py @@ -250,6 +250,11 @@ def apply_touch_verdict(self, h: int, resident: bool) -> None: if h not in self.hashes: return if not resident: + # Confirmed after the latest touch was sent (its store landed since, + # or the hash was forgotten and stored again): the touch looked + # before those bytes were there, so its verdict is stale. + if self._verified_at.get(h, -1e18) > self._touch_sent_at.get(h, -1e18): + return self.touch_misses += 1 self.forget(h) return @@ -286,7 +291,9 @@ def needs_verify(self, h: int) -> bool: existed until any evidence about `h` arrives. """ defer_s = self.defer_s - if not self.refresh_hits or defer_s <= 0 or h not in self.hashes: + # `could_serve`, not `h in hashes`: a store-only role (`kv_producer`) + # indexes its own stores but never loads them. + if not self.refresh_hits or defer_s <= 0 or not self.could_serve(h): return False now = monotonic() verified = self._verified_at.get(h) diff --git a/tests/test_state_tier_lru_refresh.py b/tests/test_state_tier_lru_refresh.py index b1fc2a66e2..daccd2194c 100644 --- a/tests/test_state_tier_lru_refresh.py +++ b/tests/test_state_tier_lru_refresh.py @@ -293,36 +293,47 @@ def test_without_a_local_cpu_backend_a_touch_asks_presence( pool.local_cpu_backend = None with caplog.at_level(logging.WARNING, logger="atom"): codec, _ = _codec(pool, max_entries=None) - assert any("no LocalCPU backend" in r.message for r in caplog.records) + assert any("no local_cpu_backend" in r.message for r in caplog.records) pool.hot[codec.key(1)] = _Obj() assert codec.touch(1) is TouchResult.PRESENT pool.hot.clear() assert codec.touch(1) is TouchResult.GONE + def _assert_degrades(self, pool, missing, caplog): + """Bind warns once naming what is missing; touches then only confirm + presence, and each one is counted rather than passing as a refresh.""" + with caplog.at_level(logging.WARNING, logger="atom"): + codec, _ = _codec(pool, max_entries=None) + bind = [r.message for r in caplog.records if missing in r.message] + assert len(bind) == 1 and "state_offload_touches_unrefreshed" in bind[0] + assert codec.put(1, []) + assert codec.touch(1) is TouchResult.PRESENT + assert codec.touch(1) is TouchResult.PRESENT + assert codec._unrefreshed_touches == 2 + @pytest.mark.parametrize("dropped", ["cpu_lock", "hot_cache", "cache_policy"]) - def test_a_local_cpu_without_the_lru_internals_fails_at_bind( - self, lmcache_key, dropped + def test_a_local_cpu_without_the_lru_internals_degrades_to_presence( + self, lmcache_key, dropped, caplog ): - """An LMCache whose names moved must not degrade every touch to a - presence check that keeps every load "succeeding".""" + """An LMCache whose internals moved keeps serving loads; the lost + refresh shows in the stats instead of taking the engine down.""" pool = _LocalCPU() setattr(pool.local_cpu_backend, dropped, None) - with pytest.raises(RuntimeError, match=dropped): - _codec(pool, max_entries=None) + self._assert_degrades(pool, f"local_cpu_backend.{dropped}", caplog) - def test_a_policy_without_update_on_hit_fails_at_bind(self, lmcache_key): + def test_a_policy_without_update_on_hit_degrades_to_presence( + self, lmcache_key, caplog + ): pool = _LocalCPU() pool.local_cpu_backend.cache_policy = SimpleNamespace() - with pytest.raises(RuntimeError, match="update_on_hit"): - _codec(pool, max_entries=None) + self._assert_degrades(pool, "cache_policy.update_on_hit", caplog) - def test_a_storage_manager_without_local_cpu_backend_fails_at_bind( - self, lmcache_key + def test_a_storage_manager_without_local_cpu_backend_degrades_to_presence( + self, lmcache_key, caplog ): pool = _LocalCPU() del pool.local_cpu_backend - with pytest.raises(RuntimeError, match="local_cpu_backend"): - _codec(pool, max_entries=None) + self._assert_degrades(pool, "no local_cpu_backend", caplog) def test_the_lru_is_resolved_once_at_bind(self, lmcache_key): pool = _LocalCPU() @@ -420,11 +431,48 @@ def test_a_hash_voted_every_pass_is_touched_at_an_interval(self, clock): def test_a_gone_verdict_drops_the_hash(self, clock): idx = _index(7) + clock[0] += 1 + idx.want_touch(7) + assert idx.take_touches() == [7] idx.apply_touch_verdict(7, False) assert not idx.could_serve(7) assert idx.touch_misses == 1 assert idx.stats()["touch_misses"] == 1 + def test_a_gone_verdict_older_than_a_store_is_ignored(self, clock): + """The touch looked before the store's bytes landed; its verdict must + not erase the image that store just indexed.""" + idx = _index(7) + clock[0] += 1 + idx.want_touch(7) + assert idx.take_touches() == [7] + clock[0] += 0.01 + idx.note_stored(7) + idx.apply_touch_verdict(7, False) + assert idx.could_serve(7) + assert idx.touch_misses == 0 + + def test_a_gone_verdict_for_a_forgotten_then_restored_hash_is_ignored(self, clock): + idx = _index(7) + clock[0] += 1 + idx.want_touch(7) + assert idx.take_touches() == [7] + idx.forget(7) + clock[0] += 0.01 + idx.note_stored(7) + idx.apply_touch_verdict(7, False) + assert idx.could_serve(7) + + def test_a_gone_verdict_after_a_store_and_a_newer_touch_applies(self, clock): + idx = _index(7) + clock[0] += 1 + idx.note_stored(7) + clock[0] += 1 + idx.want_touch(7) + assert idx.take_touches() == [7] + idx.apply_touch_verdict(7, False) + assert not idx.could_serve(7) + def test_a_resident_verdict_refreshes_the_cap_order(self, clock): idx = StateOffloadIndex(max_cpu_entries=19, refresh_hits=True) # cap 3 for h in (1, 2, 3): @@ -451,6 +499,13 @@ class TestNeedsVerify: def test_a_fresh_store_needs_no_verdict(self, clock): assert not _index(7).needs_verify(7) + def test_a_store_only_role_holds_nothing(self, clock): + """A `kv_producer` indexes its own stores but never loads them.""" + idx = _index(7, can_load=False) + clock[0] += STALE + assert not idx.needs_verify(7) + assert idx.take_touches() == [] + def test_a_stale_hash_waits_for_its_touch(self, clock): idx = _index(7) clock[0] += STALE @@ -803,6 +858,10 @@ def test_a_connector_without_touches_turns_the_refresh_off(self, clock): def test_verdicts_are_applied_to_the_index(self, clock): idx = _index(7, 8) + clock[0] += 1 + idx.want_touch(7) + idx.want_touch(8) + assert sorted(idx.take_touches()) == [7, 8] s = self._sched( SimpleNamespace(take_state_touch_verdicts=lambda: {7: False, 8: True}), idx, From 77d081e830be05b42a0e022fa43bf454b4147d6c Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Sun, 1 Nov 2026 16:35:03 +0000 Subject: [PATCH 09/18] state offload: keep LMCache reads out of BlockManager BlockManager no longer reads LMCache's CPU pool size: StateOffloadIndex takes the confirmation window (verified_for_s) instead of pool_bytes, and state_offload.state_tier_verified_for_s derives it from LMCACHE_MAX_LOCAL_CPU_SIZE, next to state_tier_chunk_tokens, which is restored as on main. Both share one LMCache config reader. Same values, same gating (only when a tier is hosted). --- atom/model_engine/block_manager.py | 16 +++--- atom/model_engine/state_offload.py | 77 +++++++++++++++++++--------- tests/test_state_tier_lru_refresh.py | 18 ++++++- 3 files changed, 75 insertions(+), 36 deletions(-) diff --git a/atom/model_engine/block_manager.py b/atom/model_engine/block_manager.py index 810b5e496e..08e5024ecf 100644 --- a/atom/model_engine/block_manager.py +++ b/atom/model_engine/block_manager.py @@ -335,8 +335,9 @@ def __init__( # releases holds a whole image out of the pool forever. from atom.model_engine.state_offload import ( state_tier_capability, + state_tier_chunk_tokens, state_tier_cpu_bytes, - state_tier_lmcache_sizes, + state_tier_verified_for_s, ) # A capability derived from the whole config, not the connector's name. @@ -461,16 +462,13 @@ def __init__( # The LMCache chunk size in tokens, read where the config is rather than # off the connector object. self._joint_chunk_tokens = 0 - lmcache_pool_bytes = 0 # Only when a joint KV load is possible at all. Reading it reaches into # LMCache, an optional dependency: probing on an engine that hosts no # state tier logged a full `ModuleNotFoundError` traceback at WARNING # on every start, for a number that engine has no use for. The # capability above is the same gate the tier itself uses. if kv_offload_enabled: - self._joint_chunk_tokens, lmcache_pool_bytes = state_tier_lmcache_sizes( - config - ) + self._joint_chunk_tokens = state_tier_chunk_tokens(config) self.state_offload: StateOffloadIndex | None = None # (req_id, hash, target_group) admitted this pass and not yet handed to # the connector. Kept here rather than in the index because the slot @@ -521,8 +519,8 @@ def __init__( # Off again (`stop_refreshing`) if the connector turns out not # to carry state touches. refresh_hits=True, - # Sizes how long a confirmed image is trusted. - pool_bytes=lmcache_pool_bytes, + # How long a confirmed image is trusted. + verified_for_s=state_tier_verified_for_s(config), ) # Attached rather than passed at construction: the coordinator is # built before the switch is read (it needs `checkpoint_spec`, @@ -1018,7 +1016,7 @@ def _joint_kv_boundary( if not seq.has_per_req_cache or self.paged_state_checkpoints is None: return self._no_joint("no_paged_checkpoints") hbs = self._hash_block_size() - # Resolved once from config at construction (`state_tier_lmcache_sizes`), + # Resolved once from config at construction (`state_tier_chunk_tokens`), # not off the connector object: the scheduler holds whatever # `get_kvconnector` returned, and one without `chunk_size` would zero # this and disable the feature silently. @@ -1555,7 +1553,7 @@ def record_allocation( With `defer_unverified`, returns False and commits nothing when the state this admission would load from the CPU tier has not been - confirmed in LMCache lately (`StateOffloadIndex.needs_verify`). The + confirmed there lately (`StateOffloadIndex.needs_verify`). The caller holds the request for a pass while that image's touch comes back, rather than parking a joint load that may miss and recompute the whole prompt. diff --git a/atom/model_engine/state_offload.py b/atom/model_engine/state_offload.py index 92ba77638a..2ed7eee0d4 100644 --- a/atom/model_engine/state_offload.py +++ b/atom/model_engine/state_offload.py @@ -58,7 +58,7 @@ def __init__( can_load: bool = True, max_cpu_entries: int | None = None, refresh_hits: bool = False, - pool_bytes: int = 0, + verified_for_s: float | None = None, defer_s: float | None = None, ) -> None: # The two legs are separately granted. A `kv_producer` role saves and @@ -116,8 +116,11 @@ def __init__( # only grow. self.refresh_hits = bool(refresh_hits) # Both fixed for the server's lifetime; `needs_verify` runs on the - # admission path. `pool_bytes` is one worker's LMCache CPU pool. - self.verified_for_s = tier_verified_for_s(pool_bytes) + # admission path. The window comes from the tier's pool + # (`state_tier_verified_for_s`); the floor when not given. + self.verified_for_s = ( + tier_verified_for_s(0) if verified_for_s is None else float(verified_for_s) + ) if defer_s is None: defer_s = verify_defer_s() if self.refresh_hits else 0.0 self.defer_s = float(defer_s) @@ -676,10 +679,29 @@ def state_tier_budget_warning( return None -def state_tier_lmcache_sizes(config) -> tuple[int, int]: - """The LMCache chunk size in tokens for the joint KV leg, and one worker's - CPU pool in bytes (`LMCACHE_MAX_LOCAL_CPU_SIZE`) for `tier_verified_for_s`; - (0, 0) when LMCache's config cannot be read. +def _lmcache_engine_config(config): + """The `LMCacheEngineConfig` this engine's offload connector builds. + + Raises when LMCache or its config cannot be read; the callers decide what + that costs. + """ + from atom.kv_transfer.offload.config import build_lmcache_config + + kvcfg = getattr(config, "kv_transfer_config", None) or {} + # Under `multi` the lmcache.* keys (chunk_size, offload_layout) live on + # the offload SUB-connector, not the composite. Passing the raw + # composite here read a zero chunk grid. Unwrap the sub the same way + # `state_tier_capability` does, so the gate's chunk size and the tier + # the capability check builds stay consistent. + if isinstance(kvcfg, dict) and kvcfg.get("kv_connector") == "multi": + sub, _why = _offload_subconfig(kvcfg) + if sub is not None: + kvcfg = sub + return build_lmcache_config(kvcfg) + + +def state_tier_chunk_tokens(config) -> int: + """The LMCache chunk size in tokens for the joint KV leg, or 0. Decided from config for the same reason `state_tier_capability` is: this runs in the engine process at `BlockManager.__init__`, and the scheduler @@ -696,21 +718,7 @@ def state_tier_lmcache_sizes(config) -> tuple[int, int]: disable the joint KV load. """ try: - from atom.kv_transfer.offload.config import build_lmcache_config - - kvcfg = getattr(config, "kv_transfer_config", None) or {} - # Under `multi` the lmcache.* keys (chunk_size, offload_layout) live on - # the offload SUB-connector, not the composite. Passing the raw - # composite here read a zero chunk grid. Unwrap the sub the same way - # `state_tier_capability` does, so the gate's chunk size and the tier - # the capability check builds stay consistent. - if isinstance(kvcfg, dict) and kvcfg.get("kv_connector") == "multi": - sub, _why = _offload_subconfig(kvcfg) - if sub is not None: - kvcfg = sub - cfg = build_lmcache_config(kvcfg) - pool_gib = float(getattr(cfg, "max_local_cpu_size", 0) or 0) - return int(cfg.chunk_size), int(pool_gib * (1 << 30)) + return int(_lmcache_engine_config(config).chunk_size) except Exception: # Blind on purpose: this runs at model load, the import reaches a # third-party package that may be absent entirely, and the only @@ -718,11 +726,30 @@ def state_tier_lmcache_sizes(config) -> tuple[int, int]: # stays off. Refusing to start would be worse. logger.warning( "state offload: could not read the LMCache chunk size; the joint " - "KV load needs it and stays off, and the CPU pool size is taken " - "as unknown", + "KV load needs it and stays off", exc_info=True, ) - return 0, 0 + return 0 + + +def state_tier_verified_for_s(config) -> float: + """`tier_verified_for_s` for one worker's LMCache CPU pool + (`LMCACHE_MAX_LOCAL_CPU_SIZE`): how long `StateOffloadIndex` trusts a + confirmed image. The floor when LMCache's config cannot be read. + + Like `state_tier_chunk_tokens`, call it only when a tier is hosted. + """ + try: + cfg = _lmcache_engine_config(config) + pool_gib = float(getattr(cfg, "max_local_cpu_size", 0) or 0) + except Exception: # noqa: BLE001 - `state_tier_chunk_tokens` logged why + logger.warning( + "state offload: LMCache's CPU pool size is unknown; a confirmed " + "state image is trusted for %.0f s", + _TIER_VERIFIED_MIN_S, + ) + pool_gib = 0.0 + return tier_verified_for_s(int(pool_gib * (1 << 30))) #: The connector backends whose worker half can build a `StateOffloadTier`. diff --git a/tests/test_state_tier_lru_refresh.py b/tests/test_state_tier_lru_refresh.py index daccd2194c..82c0e18f46 100644 --- a/tests/test_state_tier_lru_refresh.py +++ b/tests/test_state_tier_lru_refresh.py @@ -624,7 +624,7 @@ def test_an_unknown_pool_gets_the_floor(self): assert so.tier_verified_for_s(0) == so._TIER_VERIFIED_MIN_S def test_the_index_holds_by_its_own_window(self, clock): - idx = _index(7, pool_bytes=32 * self.GIB, defer_s=2) + idx = _index(7, verified_for_s=so.tier_verified_for_s(32 * self.GIB), defer_s=2) assert idx.verified_for_s == pytest.approx(8.0) clock[0] += 7.5 assert not idx.needs_verify(7) @@ -637,12 +637,26 @@ def test_the_block_manager_sizes_the_window_from_the_lmcache_pool( from test_block_manager import _OFFLOAD_KVC, _bm_with_state_tier monkeypatch.setattr( - so, "state_tier_lmcache_sizes", lambda config: (256, 100 * self.GIB) + so, + "_lmcache_engine_config", + lambda config: SimpleNamespace(chunk_size=256, max_local_cpu_size=100), ) bm = _bm_with_state_tier(monkeypatch, _OFFLOAD_KVC) assert bm._joint_chunk_tokens == 256 assert bm.state_offload.verified_for_s == pytest.approx(25.0) + def test_an_unreadable_lmcache_config_gets_the_floor(self, monkeypatch, caplog): + import sys + + from test_block_manager import _OFFLOAD_KVC, _bm_with_state_tier + + monkeypatch.setitem(sys.modules, "lmcache.v1.config", None) + with caplog.at_level(logging.WARNING, logger="atom"): + bm = _bm_with_state_tier(monkeypatch, _OFFLOAD_KVC) + assert bm._joint_chunk_tokens == 0 + assert bm.state_offload.verified_for_s == so._TIER_VERIFIED_MIN_S + assert sum(r.exc_info is not None for r in caplog.records) == 1 + # ── the connector: touches out in the metadata, verdicts back by quorum ──── From 8e9e0c9bddf80d886095d7fceeebdb054fd9a3e1 Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Mon, 2 Nov 2026 02:50:27 +0000 Subject: [PATCH 10/18] state offload: count a raising touch as unrefreshed and log it once A touch whose codec call raised was reported as refreshed, with a traceback per touch: a drifted LMCache flooded the log while state_offload_touches_unrefreshed stayed at 0. It now keeps the image indexed (resident), reports it unrefreshed, and only the first failure logs a traceback. Also name where a disk-only image goes back into LocalCPU (LMCache's StorageManager.get write-back), and fix the test fake's comment, which claimed a read-back it does not model. --- atom/kv_transfer/offload/README.md | 3 ++- .../offload/hybrid/kimi_k3/connector.py | 22 +++++++++++------ .../offload/hybrid/kimi_k3/state_object.py | 6 +++-- tests/test_state_tier_lru_refresh.py | 24 ++++++++++++++++--- 4 files changed, 42 insertions(+), 13 deletions(-) diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index d687674fc2..d50b9c3277 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -133,7 +133,8 @@ 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 +next store of it puts it back into the CPU tier. A touch that raises keeps the +image indexed and is counted there too (one traceback per worker). 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 diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py index 751d8b785d..647c41ea3d 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py @@ -134,6 +134,8 @@ def __init__(self, config) -> None: self._store_failed_no_tier: set[StateStoreOperationId] = set() # Touch verdicts since the last `get_finished`: (resident, refreshed). self._touch_verdicts: dict[StateTouchId, tuple[bool, bool]] = {} + # Touches whose codec call raised; only the first is logged. + self._touch_failures = 0 def close(self) -> None: """Drain then join the state tier before the base executors. @@ -331,13 +333,19 @@ def _touch_states(self, metadata) -> None: try: result = tier.codec.touch(int(touch_id.prefix_hash)) except Exception: - # Says nothing about the bytes; do not let it erase the index. - logger.warning( - "kimi_k3 offload: touch of state hash %d failed", - touch_id.prefix_hash, - exc_info=True, - ) - result = TouchResult.REFRESHED + # Says nothing about the bytes, so the index keeps the image; + # nothing was refreshed either. One traceback: a drifted + # LMCache raises on every touch. + self._touch_failures += 1 + if self._touch_failures == 1: + logger.warning( + "kimi_k3 offload: touch of state hash %d failed; failed " + "touches keep the image indexed and are counted in " + "state_offload_touches_unrefreshed", + touch_id.prefix_hash, + exc_info=True, + ) + result = TouchResult.PRESENT self._touch_verdicts[touch_id] = ( bool(result), result is not TouchResult.PRESENT, diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py index 4ae1dc39ae..e907f6850e 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py @@ -26,8 +26,10 @@ class TouchResult(enum.IntEnum): """What a touch found. Truthy when the image can still be read.""" GONE = 0 - #: Outside LocalCPU's hot cache but held by another backend (`get` reads - #: it and promotes it back): present, its eviction order not refreshed. + #: Outside LocalCPU's hot cache but held by another backend: present, its + #: eviction order not refreshed. `get` still reads it, through LMCache's + #: `StorageManager.get`, which writes a hit from another backend back into + #: LocalCPU (`LocalCPUBackend.submit_put_task`); nothing in ATOM does. PRESENT = 1 REFRESHED = 2 diff --git a/tests/test_state_tier_lru_refresh.py b/tests/test_state_tier_lru_refresh.py index 82c0e18f46..2aded0a403 100644 --- a/tests/test_state_tier_lru_refresh.py +++ b/tests/test_state_tier_lru_refresh.py @@ -103,8 +103,8 @@ class _LocalCPU: def __init__(self, refuse_all: bool = False) -> None: self.hot: OrderedDict = OrderedDict() - # A second backend (LocalDisk): read by `get`, which promotes a hit - # back into the hot cache, and answered by `contains`. + # A second backend (LocalDisk) for presence only: `contains` answers + # from it, `get` does not read it (no test loads a disk-only image). self.disk: dict = {} self.refuse_all = refuse_all self.fail_put = False @@ -669,6 +669,7 @@ def _k3_worker(tier): c = KimiK3OffloadConnector.__new__(KimiK3OffloadConnector) c._state_tier = tier c._touch_verdicts = {} + c._touch_failures = 0 return c @@ -705,7 +706,24 @@ def touch(h): worker = _k3_worker(SimpleNamespace(codec=SimpleNamespace(touch=touch))) worker._touch_states(SimpleNamespace(state_touches=[StateTouchId(1, 0)])) - assert worker._touch_verdicts == {StateTouchId(1, 0): (True, True)} + assert worker._touch_verdicts[StateTouchId(1, 0)][0] is True + + def test_a_failing_touch_counts_unrefreshed_and_logs_once(self, caplog): + """A drifted LMCache raises on every touch: that must show up in + `touches_unrefreshed`, not as healthy refreshes, and not as one + traceback per touch.""" + + def touch(h): + raise RuntimeError("backend hiccup") + + worker = _k3_worker(SimpleNamespace(codec=SimpleNamespace(touch=touch))) + ids = [StateTouchId(1, 0), StateTouchId(2, 1), StateTouchId(1, 2)] + with caplog.at_level(logging.WARNING, logger="atom"): + worker._touch_states(SimpleNamespace(state_touches=ids[:2])) + worker._touch_states(SimpleNamespace(state_touches=ids[2:])) + assert worker._touch_verdicts == {i: (True, False) for i in ids} + assert worker._touch_failures == 3 + assert [r.exc_info is not None for r in caplog.records] == [True] def _scheduler(self, monkeypatch): from atom.kv_transfer.offload.dense.connector import DenseOffloadScheduler From 0dda9c9d3fce0069991fd5788df1ed881c1fe1e8 Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Mon, 2 Nov 2026 03:35:28 +0000 Subject: [PATCH 11/18] state offload: keep the tier logic in state_offload, tests with their components The block manager and the scheduler now only call the index: - `StateOffloadIndex.hold_admission` makes the hold decision and counts it; the block manager's `_resume_hash_off_hbm` only says which image an admission would load from off-HBM. - `request_load` registers its own touch, so `_request_state_load` no longer does. - `publish_touches` / `collect_touch_verdicts` replace the scheduler's `_publish_state_touches` / `_apply_state_touch_verdicts`. - `_reachable` calls `want_touch` directly (the getattr shim for an index without touches is gone; every index has them). tests/test_state_tier_lru_refresh.py is split into the existing files: codec and connector touches -> test_lmcache_offload_connector.py, index and budget -> test_state_offload_index.py, window sizing -> test_block_manager.py, admission hold and voters -> test_state_checkpoint.py. --- atom/kv_transfer/disaggregation/types.py | 3 +- atom/model_engine/block_manager.py | 54 +- atom/model_engine/page_unit_checkpoint.py | 7 +- atom/model_engine/scheduler.py | 50 +- atom/model_engine/state_offload.py | 54 +- tests/test_block_manager.py | 28 + tests/test_lmcache_offload_connector.py | 521 +++++++++ tests/test_state_checkpoint.py | 293 +++++ tests/test_state_offload_index.py | 355 +++++- tests/test_state_tier_lru_refresh.py | 1224 --------------------- 10 files changed, 1283 insertions(+), 1306 deletions(-) delete mode 100644 tests/test_state_tier_lru_refresh.py diff --git a/atom/kv_transfer/disaggregation/types.py b/atom/kv_transfer/disaggregation/types.py index d2bce7bf08..b46edb1036 100644 --- a/atom/kv_transfer/disaggregation/types.py +++ b/atom/kv_transfer/disaggregation/types.py @@ -112,7 +112,8 @@ class StateStoreOperationId: 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. + settled and its bytes are never re-indexed. `generation` is what separates + the attempts. """ prefix_hash: int diff --git a/atom/model_engine/block_manager.py b/atom/model_engine/block_manager.py index 08e5024ecf..8ed11fc226 100644 --- a/atom/model_engine/block_manager.py +++ b/atom/model_engine/block_manager.py @@ -1552,11 +1552,9 @@ def record_allocation( """Commit the fit probe before any allocation changes the pool. With `defer_unverified`, returns False and commits nothing when the - state this admission would load from the CPU tier has not been - confirmed there lately (`StateOffloadIndex.needs_verify`). The - caller holds the request for a pass while that image's touch comes - back, rather than parking a joint load that may miss and recompute the - whole prompt. + state offload asks to hold the admission over the state it would + restore from outside HBM (`StateOffloadIndex.hold_admission`). The + caller retries the request a pass later. """ if not self.enable_prefix_caching: return True @@ -1564,34 +1562,37 @@ def record_allocation( seq, probe.num_cached_blocks, probe.block_hashes ) decision = self._fit_joint(decision, probe) - if defer_unverified and self._tier_state_unverified( - seq, probe.num_cached_blocks, probe.block_hashes, decision + if ( + defer_unverified + and self.state_offload is not None + and self.state_offload.hold_admission( + self._resume_hash_off_hbm( + seq, probe.num_cached_blocks, probe.block_hashes, decision + ) + ) ): - self.state_offload.loads_deferred_for_verify += 1 return False self._commit_joint_boundary(seq, decision) return True - def _tier_state_unverified( + def _resume_hash_off_hbm( self, seq: Sequence, num_cached_blocks: int, block_hashes: list[int], decision: JointBoundaryDecision, - ) -> bool: - """Whether `allocate` would load this seq's state from the CPU tier - under a hash not confirmed resident recently. + ) -> int: + """The hash `allocate` would restore this seq's state from when HBM + does not hold its state, or -1 (no state, HBM holds it, or a replay). - Judges the hash `allocate` will hand `_attach_state_slots` - (`_state_resume_hash`), which the tier only serves when HBM does not - hold it. + The hash `allocate` will hand `_attach_state_slots` + (`_state_resume_hash`), so a hold judges what will be loaded. """ - offload = self.state_offload - if offload is None or not offload.refresh_hits or not seq.has_per_req_cache: - return False + if not seq.has_per_req_cache: + return -1 joint = decision.skip_reason is None if joint and decision.state_from_hbm: - return False + return -1 hit_hash = ( block_hashes[num_cached_blocks - 1] if 0 < num_cached_blocks <= len(block_hashes) @@ -1600,15 +1601,13 @@ def _tier_state_unverified( h = self._state_resume_hash( seq, decision.boundary_hash if joint else -1, hit_hash ) - if h == -1: - return False - if not joint: + if h != -1 and not joint: if self.paged_state_checkpoints is not None: if self.paged_state_checkpoints.contains(h): - return False + return -1 elif self.state.lookup(h) >= 0: - return False - return offload.needs_verify(h) + return -1 + return h def _state_resume_hash(self, seq: Sequence, joint_hash: int, hit_hash: int) -> int: """The hash whose state `seq` resumes from, or -1 for none. @@ -1616,8 +1615,8 @@ def _state_resume_hash(self, seq: Sequence, joint_hash: int, hit_hash: int) -> i The joint boundary when there is one (above the HBM hit by construction), else the HBM hit's last block. A replay rebuilds the state instead of restoring it, so it resumes from none. `allocate` - hands this to `_attach_state_slots`, and `_tier_state_unverified` asks - about it, so the admission hold judges the hash that will be loaded. + hands this to `_attach_state_slots`, and `_resume_hash_off_hbm` reports + it, so the admission hold judges the hash that will be loaded. """ if joint_hash != -1: return joint_hash @@ -2116,7 +2115,6 @@ def _request_state_load(self, seq: Sequence, hit_hash: int) -> bool: return False if not self.state_offload.request_load(seq.id, hit_hash): return False - self.state_offload.want_touch(hit_hash) seq.offload_joint.load_hash = hit_hash self._state_loads.append((seq.id, hit_hash, seq.state_slot)) return True diff --git a/atom/model_engine/page_unit_checkpoint.py b/atom/model_engine/page_unit_checkpoint.py index 00051ffedd..2adb966c81 100644 --- a/atom/model_engine/page_unit_checkpoint.py +++ b/atom/model_engine/page_unit_checkpoint.py @@ -1080,12 +1080,7 @@ def _reachable(self, h: int) -> bool: return True if self.offload is None or not self.offload.could_serve(h): return False - # `getattr`, as for `stats`: `attach_offload` accepts anything that - # answers `could_serve`, and one without touches loses the LRU refresh, - # not the hit. - want_touch = getattr(self.offload, "want_touch", None) - if callable(want_touch): - want_touch(h) + self.offload.want_touch(h) return True def checkpoint(self, seq: Sequence, boundary_blocks: int, h: int) -> None: diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index e49b5b8e48..8965792dd2 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -2382,10 +2382,10 @@ def _schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]] | None: if not self.block_manager.record_allocation( seq, hbm_probe, defer_unverified=True ): - # The CPU state image this admission would load has not been - # confirmed in LMCache lately and its touch is out. Held like - # an in-flight prefix: an image found gone is forgotten when - # the verdict lands, and the retry plans without it. + # The state offload holds it a pass (the state it would load + # is unconfirmed, and its check is out). Held like an + # in-flight prefix: an image found gone is forgotten when the + # verdict lands, and the retry plans without it. prefix_waiters.append(seq) continue if not self.block_manager.allocate(seq, num_cached_blocks): @@ -3779,42 +3779,6 @@ def _publish_state_stores(self) -> None: # never reached a worker, so it is not a `stores_failed`. self.block_manager.settle_state_store(op, ok=False, attempted=False) - def _publish_state_touches(self) -> None: - """Hand the state images this pass voted a hit for to the connector, - which refreshes them in LMCache's LRU and reports whether each is - still there (`_apply_state_touch_verdicts`). - - A connector that will not carry them turns the refresh off for good: - otherwise every admission that waits for a verdict would wait out - `OFFLOAD_LOOKUP_DEFER_S` for nothing. - """ - offload = getattr(getattr(self, "block_manager", None), "state_offload", None) - if offload is None or not offload.refresh_hits: - return - hashes = offload.take_touches() - if not hashes: - return - enqueue = getattr(self.kv_connector, "enqueue_state_touches", None) - if enqueue is not None and enqueue(hashes): - return - logger.warning( - "state offload: %s does not carry state touches; state images " - "will not be refreshed in LMCache's LRU, and an image LMCache " - "evicts is only noticed when a load misses.", - type(self.kv_connector).__name__, - ) - offload.stop_refreshing() - - def _apply_state_touch_verdicts(self, offload) -> None: - take = getattr(self.kv_connector, "take_state_touch_verdicts", None) - if take is None: - return - for h, resident in take().items(): - offload.apply_touch_verdict(h, resident) - unrefreshed = getattr(self.kv_connector, "take_state_touches_unrefreshed", None) - if unrefreshed is not None: - offload.touches_unrefreshed += int(unrefreshed()) - def _publish_state_loads(self) -> None: """Hand this pass's state-tier loads to the connector, before it builds its metadata. @@ -3825,10 +3789,12 @@ def _publish_state_loads(self) -> None: """ if self.kv_connector is None: return - self._publish_state_touches() # Guarded like the `_settle_state_load` / `_publish_state_stores` # siblings: a scheduler built without a block manager (the connector # test doubles) would otherwise AttributeError on the bare deref. + offload = getattr(getattr(self, "block_manager", None), "state_offload", None) + if offload is not None: + offload.publish_touches(self.kv_connector) take = getattr(getattr(self, "block_manager", None), "take_state_loads", None) if take is None: return @@ -4953,7 +4919,7 @@ def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput): offload = getattr(bm, "state_offload", None) if offload is not None: # Before the next admission pass, which reads what they decide. - self._apply_state_touch_verdicts(offload) + offload.collect_touch_verdicts(self.kv_connector) take = getattr(self.kv_connector, "take_state_reports", None) if offload is not None and take is not None: indexed, failed = take() diff --git a/atom/model_engine/state_offload.py b/atom/model_engine/state_offload.py index 2ed7eee0d4..e27dc4e23c 100644 --- a/atom/model_engine/state_offload.py +++ b/atom/model_engine/state_offload.py @@ -215,8 +215,8 @@ def want_touch(self, h: int) -> None: """A voter counted `h` as a hit: under `refresh_hits`, queue it for `take_touches`. - Called by every voter on the hits it votes for (the resumable scan and - `BlockManager._request_state_load`), so a prefix whose state is voted + Called on every hit voted for (the resumable scan, and `request_load` + for the load it offers), so a prefix whose state is voted for on every pass stays as warm in LMCache as its KV chunks, which the KV lookup refreshes on every pass. """ @@ -274,6 +274,43 @@ def stop_refreshing(self) -> None: self._hold_started_at.clear() self._hold_expired.clear() + def publish_touches(self, connector) -> None: + """Hand this pass's touches to `connector`, which refreshes the images + in LMCache's LRU and reports whether each is still there + (`collect_touch_verdicts`). + + A connector that will not carry them turns the refresh off for good: + otherwise every admission that waits for a verdict would wait out + `defer_s` for nothing. + """ + if not self.refresh_hits: + return + hashes = self.take_touches() + if not hashes: + return + enqueue = getattr(connector, "enqueue_state_touches", None) + if enqueue is not None and enqueue(hashes): + return + logger.warning( + "state offload: %s does not carry state touches; state images " + "will not be refreshed in LMCache's LRU, and an image LMCache " + "evicts is only noticed when a load misses.", + type(connector).__name__, + ) + self.stop_refreshing() + + def collect_touch_verdicts(self, connector) -> None: + """Apply the touch verdicts `connector` reached TP quorum on, and add + its presence-only touches to `touches_unrefreshed`.""" + take = getattr(connector, "take_state_touch_verdicts", None) + if take is None: + return + for h, resident in take().items(): + self.apply_touch_verdict(h, resident) + unrefreshed = getattr(connector, "take_state_touches_unrefreshed", None) + if unrefreshed is not None: + self.touches_unrefreshed += int(unrefreshed()) + def needs_verify(self, h: int) -> bool: """Whether an admission about to load `h` should wait a pass for a touch verdict instead of parking on it now. @@ -321,13 +358,23 @@ def needs_verify(self, h: int) -> bool: self._touch_wanted.add(h) return True + def hold_admission(self, h: int) -> bool: + """Whether an admission that would restore its state from `h` (-1: + from nothing outside HBM) waits a pass for `h`'s touch verdict + (`needs_verify`). Counted in `loads_deferred_for_verify`.""" + if h == -1 or not self.needs_verify(h): + return False + self.loads_deferred_for_verify += 1 + return True + # -------------------------------- loads -------------------------------- # def request_load(self, req_id, h: int) -> bool: """Offer to fetch `h` back for `req_id`. False if this tier cannot. The guard between believing and delivering: a load is resolved only by a worker report, so offering one for a hash never stored would park the - request against bytes no `get` can produce. + request against bytes no `get` can produce. An offered load is a hit, + so it is registered for a touch (`want_touch`). """ if not self.could_serve(h): # A store-only role (can_load False) or a hash never stored. Voting @@ -347,6 +394,7 @@ def request_load(self, req_id, h: int) -> bool: return False self.pending_loads[req_id] = int(h) self.loads_attempted += 1 + self.want_touch(h) return True def complete_load(self, req_id) -> None: diff --git a/tests/test_block_manager.py b/tests/test_block_manager.py index 0d29539c1c..7533a9a4cb 100644 --- a/tests/test_block_manager.py +++ b/tests/test_block_manager.py @@ -2,11 +2,14 @@ # Tests for atom/model_engine/block_manager.py — public API only import logging +import sys +from types import SimpleNamespace import pytest from conftest import MockConfig from atom.model_engine import block_manager as bm_module +from atom.model_engine import state_offload as so from atom.model_engine.block_manager import BlockManager # ── compute_hash ─────────────────────────────────────────────────────────── @@ -1368,3 +1371,28 @@ def test_a_hosted_tier_still_probes_and_says_so_when_it_cannot_read( # Exercise the missing dependency even in images with LMCache installed. assert bm._joint_chunk_tokens == 0 assert "LMCache chunk size" in caplog.text + + +class TestStateTierVerifyWindow: + """The state tier's confirmation window comes from the LMCache pool, + read in `state_offload`, not here.""" + + def test_the_block_manager_sizes_the_window_from_the_lmcache_pool( + self, monkeypatch + ): + monkeypatch.setattr( + so, + "_lmcache_engine_config", + lambda config: SimpleNamespace(chunk_size=256, max_local_cpu_size=100), + ) + bm = _bm_with_state_tier(monkeypatch, _OFFLOAD_KVC) + assert bm._joint_chunk_tokens == 256 + assert bm.state_offload.verified_for_s == pytest.approx(25.0) + + def test_an_unreadable_lmcache_config_gets_the_floor(self, monkeypatch, caplog): + monkeypatch.setitem(sys.modules, "lmcache.v1.config", None) + with caplog.at_level(logging.WARNING, logger="atom"): + bm = _bm_with_state_tier(monkeypatch, _OFFLOAD_KVC) + assert bm._joint_chunk_tokens == 0 + assert bm.state_offload.verified_for_s == so._TIER_VERIFIED_MIN_S + assert sum(r.exc_info is not None for r in caplog.records) == 1 diff --git a/tests/test_lmcache_offload_connector.py b/tests/test_lmcache_offload_connector.py index a6546babcf..b7124e33db 100644 --- a/tests/test_lmcache_offload_connector.py +++ b/tests/test_lmcache_offload_connector.py @@ -45,6 +45,7 @@ LoadOperationId, SaveOperationId, StateStoreOperationId, + StateTouchId, ) from atom.kv_transfer.offload import _block_gpu_connector from atom.kv_transfer.offload import config as offcfg @@ -78,6 +79,7 @@ KimiK3OffloadScheduler, save_stall_seconds, ) +from atom.kv_transfer.offload.hybrid.kimi_k3.state_object import TouchResult from atom.kv_transfer.offload.hybrid.kimi_k3.state_tier import _JointPark from atom.kv_transfer.offload.hybrid.qwen.connector import QwenOffloadScheduler from atom.kv_transfer.offload.metadata import ( @@ -7346,3 +7348,522 @@ def get(self, h, slot): assert tier.take_store_reports() == (set(), {op}) finally: tier.shutdown() + + +class _Obj: + """A MemoryObj's reference count, and nothing else.""" + + def __init__(self, size: int = 1024) -> None: + self.refs = 1 + self.size = size + + def ref_count_up(self) -> None: + self.refs += 1 + + def ref_count_down(self) -> None: + assert self.refs > 0, "reference discharged twice" + self.refs -= 1 + + def get_ref_count(self) -> int: + return self.refs + + def get_size(self) -> int: + return self.size + + +class _LRUPolicy: + """LMCache's `LRUCachePolicy.update_on_hit`.""" + + def update_on_hit(self, key, cache) -> None: + cache.move_to_end(key) + + +class _LocalCPU: + """`storage_manager` over LocalCPU's rules. + + `batched_put` adopts an object with its own reference unless the key is + already resident (LocalCPU refuses duplicates), then drops the caller's. + `get` does not touch the LRU (`get_blocking` only takes a reference). + `evict_oldest` is the pool making room for KV: oldest first, skipping + anything another holder still references. + """ + + def __init__(self, refuse_all: bool = False) -> None: + self.hot: OrderedDict = OrderedDict() + # A second backend (LocalDisk) for presence only: `contains` answers + # from it, `get` does not read it (no test loads a disk-only image). + self.disk: dict = {} + self.refuse_all = refuse_all + self.fail_put = False + self.puts = 0 + self.local_cpu_backend = SimpleNamespace( + cpu_lock=threading.Lock(), hot_cache=self.hot, cache_policy=_LRUPolicy() + ) + + def batched_put(self, keys, objs) -> None: + if self.fail_put: + raise RuntimeError("put failed") + self.puts += 1 + for k, o in zip(keys, objs, strict=True): + if not self.refuse_all and k not in self.hot: + o.ref_count_up() + self.hot[k] = o + o.ref_count_down() + + def get(self, key): + o = self.hot.get(key) + if o is not None: + o.ref_count_up() + return o + + def contains(self, key): + if key in self.hot: + return "LocalCPUBackend" + return "LocalDiskBackend" if key in self.disk else None + + def spill_to_disk(self, key) -> None: + """LocalCPU evicted `key`; the disk backend still holds a copy.""" + self.disk[key] = True + self.hot.pop(key).ref_count_down() + + def remove(self, key) -> int: + o = self.hot.pop(key, None) + if o is None: + return 0 + o.ref_count_down() + return 1 + + def evict_oldest(self): + for k, o in self.hot.items(): + if o.get_ref_count() == 1: + self.hot.pop(k).ref_count_down() + return k + return None + + +def _lru_codec(storage, max_entries): + from atom.kv_transfer.offload.hybrid.kimi_k3.state_object import StateByteCodec + + codec = StateByteCodec( + backend=SimpleNamespace( + page_unit_views=lambda units: [], state_entry_views=lambda slot: [] + ), + staged=SimpleNamespace(pack=lambda views, obj: None, unpack=lambda o, v: None), + entry_bytes=1024, + model_name="qwen", + world_size=1, + worker_id=0, + layout_id="L", + max_cache_bytes=None if max_entries is None else max_entries * 1024, + ) + codec.bind_storage_manager(storage) + allocated = [] + + def allocate(nbytes): + allocated.append(_Obj(nbytes)) + return allocated[-1] + + codec._allocate = allocate + return codec, allocated + + +def _kv(pool, name): + pool.hot[name] = _Obj() + + +# ── the codec: touch, get and re-store refresh LMCache's LRU ─────────────── + + +class TestCodecRefresh: + def test_an_untouched_image_is_the_first_thing_pressure_evicts(self, lmcache_key): + """The bug: KV written after the image is younger, however often the + prefix hits.""" + pool = _LocalCPU() + codec, _ = _lru_codec(pool, max_entries=None) + assert codec.put(1, []) + _kv(pool, "kv-chunk") + assert pool.evict_oldest() == codec.key(1) + + def test_a_touch_moves_the_image_past_older_kv(self, lmcache_key): + pool = _LocalCPU() + codec, _ = _lru_codec(pool, max_entries=None) + assert codec.put(1, []) + _kv(pool, "kv-chunk") + assert codec.touch(1) is TouchResult.REFRESHED + assert pool.evict_oldest() == "kv-chunk" + assert codec.key(1) in pool.hot + + def test_a_touch_reports_an_evicted_image_gone(self, lmcache_key): + pool = _LocalCPU() + codec, _ = _lru_codec(pool, max_entries=2) + assert codec.put(1, []) + assert pool.evict_oldest() == codec.key(1) + assert codec.touch(1) is TouchResult.GONE + assert codec.key(1) not in codec._cache_keys, "no longer counts against the cap" + + def test_a_resume_refreshes_the_image(self, lmcache_key): + """`get_blocking` takes a reference and nothing else.""" + pool = _LocalCPU() + codec, objs = _lru_codec(pool, max_entries=None) + assert codec.put(1, []) + _kv(pool, "kv-chunk") + assert codec.get(1, 0) + assert objs[0].refs == 1, "the load's reference is discharged" + assert pool.evict_oldest() == "kv-chunk" + + def test_a_missed_get_leaves_the_cap_lru(self, lmcache_key): + pool = _LocalCPU() + codec, _ = _lru_codec(pool, max_entries=2) + assert codec.put(1, []) + pool.evict_oldest() + assert not codec.get(1, 0) + assert codec._cache_keys == OrderedDict() + + def test_a_resident_image_is_refreshed_not_stored_twice(self, lmcache_key): + pool = _LocalCPU() + codec, objs = _lru_codec(pool, max_entries=2) + released = [] + assert codec.put(1, []) + _kv(pool, "kv-chunk") + assert codec.put(1, [], on_source_released=lambda: released.append(True)) + + assert len(objs) == 1, "no second image-sized allocation and D2H" + assert released == [True], "the source units must still be handed back" + assert pool.evict_oldest() == "kv-chunk" + + def test_an_evicted_image_is_stored_again(self, lmcache_key): + pool = _LocalCPU() + codec, objs = _lru_codec(pool, max_entries=2) + assert codec.put(1, []) + pool.evict_oldest() + assert codec.put(1, []) + assert len(objs) == 2 and codec.key(1) in pool.hot + + def test_the_cap_retires_the_least_recently_used_image(self, lmcache_key): + """Not the oldest store: the image a touch kept warm survives.""" + pool = _LocalCPU() + codec, objs = _lru_codec(pool, max_entries=2) + assert codec.put(1, []) and codec.put(2, []) + assert codec.touch(1) is TouchResult.REFRESHED + assert codec.put(3, []) + + assert codec.key(2) not in pool.hot + assert objs[1].refs == 0, "a retired image goes back to the allocator" + assert list(codec._cache_keys) == [codec.key(1), codec.key(3)] + + def test_images_hold_no_extra_reference(self, lmcache_key): + """No reservation: KV pressure may evict an image like any chunk.""" + pool = _LocalCPU() + codec, objs = _lru_codec(pool, max_entries=2) + assert codec.put(1, []) + assert objs[0].refs == 1 + assert pool.evict_oldest() == codec.key(1) + + def test_a_refused_put_keeps_no_reference(self, lmcache_key): + pool = _LocalCPU(refuse_all=True) + codec, objs = _lru_codec(pool, max_entries=2) + assert codec.put(1, []) + assert objs[0].refs == 0, "holding a refused object would strand its block" + + def test_a_failed_put_discharges_its_reference(self, lmcache_key): + pool = _LocalCPU() + pool.fail_put = True + codec, objs = _lru_codec(pool, max_entries=2) + with pytest.raises(RuntimeError): + codec.put(1, []) + assert objs[0].refs == 0 + + def test_without_a_local_cpu_backend_a_touch_asks_presence( + self, lmcache_key, caplog + ): + pool = _LocalCPU() + pool.local_cpu_backend = None + with caplog.at_level(logging.WARNING, logger="atom"): + codec, _ = _lru_codec(pool, max_entries=None) + assert any("no local_cpu_backend" in r.message for r in caplog.records) + pool.hot[codec.key(1)] = _Obj() + assert codec.touch(1) is TouchResult.PRESENT + pool.hot.clear() + assert codec.touch(1) is TouchResult.GONE + + def _assert_degrades(self, pool, missing, caplog): + """Bind warns once naming what is missing; touches then only confirm + presence, and each one is counted rather than passing as a refresh.""" + with caplog.at_level(logging.WARNING, logger="atom"): + codec, _ = _lru_codec(pool, max_entries=None) + bind = [r.message for r in caplog.records if missing in r.message] + assert len(bind) == 1 and "state_offload_touches_unrefreshed" in bind[0] + assert codec.put(1, []) + assert codec.touch(1) is TouchResult.PRESENT + assert codec.touch(1) is TouchResult.PRESENT + assert codec._unrefreshed_touches == 2 + + @pytest.mark.parametrize("dropped", ["cpu_lock", "hot_cache", "cache_policy"]) + def test_a_local_cpu_without_the_lru_internals_degrades_to_presence( + self, lmcache_key, dropped, caplog + ): + """An LMCache whose internals moved keeps serving loads; the lost + refresh shows in the stats instead of taking the engine down.""" + pool = _LocalCPU() + setattr(pool.local_cpu_backend, dropped, None) + self._assert_degrades(pool, f"local_cpu_backend.{dropped}", caplog) + + def test_a_policy_without_update_on_hit_degrades_to_presence( + self, lmcache_key, caplog + ): + pool = _LocalCPU() + pool.local_cpu_backend.cache_policy = SimpleNamespace() + self._assert_degrades(pool, "cache_policy.update_on_hit", caplog) + + def test_a_storage_manager_without_local_cpu_backend_degrades_to_presence( + self, lmcache_key, caplog + ): + pool = _LocalCPU() + del pool.local_cpu_backend + self._assert_degrades(pool, "no local_cpu_backend", caplog) + + def test_the_lru_is_resolved_once_at_bind(self, lmcache_key): + pool = _LocalCPU() + codec, _ = _lru_codec(pool, max_entries=None) + assert codec.put(1, []) + _kv(pool, "kv-chunk") + pool.local_cpu_backend = None + assert codec.touch(1) is TouchResult.REFRESHED + assert pool.evict_oldest() == "kv-chunk" + + def test_a_disk_only_image_is_counted_and_warned_once(self, lmcache_key, caplog): + """Present, so still readable, but its LRU was not refreshed: it must + not be silent.""" + pool = _LocalCPU() + codec, _ = _lru_codec(pool, max_entries=None) + assert codec.put(1, []) + pool.spill_to_disk(codec.key(1)) + with caplog.at_level(logging.WARNING, logger="atom"): + assert codec.touch(1) is TouchResult.PRESENT + assert codec.touch(1) is TouchResult.PRESENT + assert codec._unrefreshed_touches == 2 + warned = [r for r in caplog.records if "touches_unrefreshed" in r.message] + assert len(warned) == 1 + + def test_a_refreshed_or_gone_image_is_not_counted(self, lmcache_key, caplog): + pool = _LocalCPU() + codec, _ = _lru_codec(pool, max_entries=None) + assert codec.put(1, []) + with caplog.at_level(logging.WARNING, logger="atom"): + assert codec.touch(1) is TouchResult.REFRESHED + pool.hot.clear() + assert codec.touch(1) is TouchResult.GONE + assert codec._unrefreshed_touches == 0 + assert not [r for r in caplog.records if "touches_unrefreshed" in r.message] + + def test_a_disk_only_image_is_stored_back_into_the_cpu_tier(self, lmcache_key): + """The put fast path is for an image LocalCPU still holds, not for + one only the disk does.""" + pool = _LocalCPU() + codec, objs = _lru_codec(pool, max_entries=None) + assert codec.put(1, []) + pool.spill_to_disk(codec.key(1)) + released = [] + assert codec.put(1, [], on_source_released=lambda: released.append(True)) + assert pool.puts == 2 and len(objs) == 2 + assert codec.key(1) in pool.hot + assert released == [True] + assert codec._unrefreshed_touches == 0, "a store is not a touch verdict" + + +# ── the connector: touches out in the metadata, verdicts back by quorum ──── + + +def _touch_worker(tier): + from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( + KimiK3OffloadConnector, + ) + + c = KimiK3OffloadConnector.__new__(KimiK3OffloadConnector) + c._state_tier = tier + c._touch_verdicts = {} + c._touch_failures = 0 + return c + + +class TestConnectorTouches: + def test_the_worker_touches_and_records_each_verdict(self): + touched = [] + + found = {1: TouchResult.REFRESHED, 2: TouchResult.GONE, 3: TouchResult.PRESENT} + + def touch(h): + touched.append(h) + return found[h] + + worker = _touch_worker(SimpleNamespace(codec=SimpleNamespace(touch=touch))) + ids = [StateTouchId(1, 0), StateTouchId(2, 1), StateTouchId(3, 2)] + worker._touch_states(SimpleNamespace(state_touches=ids)) + assert touched == [1, 2, 3] + # (resident, refreshed): only a present-but-unrefreshed image is False + # on the second. + assert worker._touch_verdicts == { + ids[0]: (True, True), + ids[1]: (False, True), + ids[2]: (True, False), + } + + def test_with_no_tier_every_image_is_gone(self): + worker = _touch_worker(None) + worker._touch_states(SimpleNamespace(state_touches=[StateTouchId(1, 0)])) + assert worker._touch_verdicts == {StateTouchId(1, 0): (False, True)} + + def test_a_failing_touch_does_not_erase_the_index(self): + def touch(h): + raise RuntimeError("backend hiccup") + + worker = _touch_worker(SimpleNamespace(codec=SimpleNamespace(touch=touch))) + worker._touch_states(SimpleNamespace(state_touches=[StateTouchId(1, 0)])) + assert worker._touch_verdicts[StateTouchId(1, 0)][0] is True + + def test_a_failing_touch_counts_unrefreshed_and_logs_once(self, caplog): + """A drifted LMCache raises on every touch: that must show up in + `touches_unrefreshed`, not as healthy refreshes, and not as one + traceback per touch.""" + + def touch(h): + raise RuntimeError("backend hiccup") + + worker = _touch_worker(SimpleNamespace(codec=SimpleNamespace(touch=touch))) + ids = [StateTouchId(1, 0), StateTouchId(2, 1), StateTouchId(1, 2)] + with caplog.at_level(logging.WARNING, logger="atom"): + worker._touch_states(SimpleNamespace(state_touches=ids[:2])) + worker._touch_states(SimpleNamespace(state_touches=ids[2:])) + assert worker._touch_verdicts == {i: (True, False) for i in ids} + assert worker._touch_failures == 3 + assert [r.exc_info is not None for r in caplog.records] == [True] + + def _scheduler(self, monkeypatch): + from atom.kv_transfer.offload.dense.connector import DenseOffloadScheduler + from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( + KimiK3OffloadScheduler, + ) + from atom.kv_transfer.offload.metadata import LMCacheOffloadMetadata + + monkeypatch.setattr( + DenseOffloadScheduler, + "build_connector_meta", + lambda self: LMCacheOffloadMetadata(), + ) + s = KimiK3OffloadScheduler.__new__(KimiK3OffloadScheduler) + s._pending_state_loads = [] + s._pending_state_stores = [] + s._pending_state_touches = [] + s._touch_round = 0 + s._state_touch_verdicts = {} + s._state_touches_unrefreshed = 0 + s._refresh_save_stall = lambda: None + return s + + def test_each_touch_is_a_distinct_identity_drained_once(self, monkeypatch): + s = self._scheduler(monkeypatch) + assert s.enqueue_state_touches([]) is False + assert s.enqueue_state_touches([5, 5]) + meta = s.build_connector_meta() + assert meta.state_touches == [StateTouchId(5, 0), StateTouchId(5, 1)] + assert s.build_connector_meta().state_touches == [] + + def test_one_rank_missing_the_image_makes_it_gone(self, monkeypatch): + from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( + STATE_TOUCH_CHANNEL, + ) + + s = self._scheduler(monkeypatch) + for op, ok in ((StateTouchId(5, 0), True), (StateTouchId(5, 1), False)): + handled = s.connector_completion( + ConnectorCompletion(STATE_TOUCH_CHANNEL, op, ok) + ) + assert handled is None, "handled, but not a terminal save" + s.connector_completion( + ConnectorCompletion(STATE_TOUCH_CHANNEL, StateTouchId(6, 2), True) + ) + assert s.take_state_touch_verdicts() == {5: False, 6: True} + assert s.take_state_touch_verdicts() == {} + + def test_a_presence_only_shard_counts_the_touch_unrefreshed(self, monkeypatch): + from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( + STATE_TOUCH_REFRESH_CHANNEL, + ) + + s = self._scheduler(monkeypatch) + for op, ok in ((StateTouchId(5, 0), True), (StateTouchId(6, 1), False)): + handled = s.connector_completion( + ConnectorCompletion(STATE_TOUCH_REFRESH_CHANNEL, op, ok) + ) + assert handled is None, "handled, but not a terminal save" + assert s.take_state_touches_unrefreshed() == 1 + assert s.take_state_touches_unrefreshed() == 0 + assert s.take_state_touch_verdicts() == {}, "not a residency verdict" + + def test_both_touch_reports_reach_quorum_across_tp(self, monkeypatch): + """Each rank files a residency and a refresh report per touch; one + rank that only found the image present counts the touch unrefreshed + without making it gone.""" + from atom.kv_transfer.disaggregation.aggregator import KVOutputAggregator + from atom.kv_transfer.disaggregation.types import KVConnectorOutput + from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( + STATE_TOUCH_CHANNEL, + STATE_TOUCH_REFRESH_CHANNEL, + ) + + touch = StateTouchId(5, 0) + outputs = [] + for result in (TouchResult.REFRESHED, TouchResult.PRESENT): + worker = _touch_worker( + SimpleNamespace(codec=SimpleNamespace(touch=lambda h, r=result: r)) + ) + worker._touch_states(SimpleNamespace(state_touches=[touch])) + resident, refreshed = worker._touch_verdicts[touch] + outputs.append( + KVConnectorOutput( + connector_completions={ + ConnectorCompletion(STATE_TOUCH_CHANNEL, touch, resident), + ConnectorCompletion( + STATE_TOUCH_REFRESH_CHANNEL, touch, refreshed + ), + } + ) + ) + merged = KVOutputAggregator(world_size=2).aggregate(outputs) + s = self._scheduler(monkeypatch) + for completion in merged.connector_completions: + s.connector_completion(completion) + assert s.take_state_touch_verdicts() == {5: True} + assert s.take_state_touches_unrefreshed() == 1 + + def test_the_shells_forward_the_unrefreshed_count(self, monkeypatch): + from atom.kv_transfer.disaggregation.multi.multi_connector import ( + MultiConnectorScheduler, + ) + from atom.kv_transfer.offload.connector import ( + LMCacheOffloadConnectorScheduler, + ) + + shell = LMCacheOffloadConnectorScheduler.__new__( + LMCacheOffloadConnectorScheduler + ) + shell._impl = SimpleNamespace(take_state_touches_unrefreshed=lambda: 2) + assert shell.take_state_touches_unrefreshed() == 2 + shell._impl = SimpleNamespace() + assert shell.take_state_touches_unrefreshed() == 0 + + multi = MultiConnectorScheduler.__new__(MultiConnectorScheduler) + monkeypatch.setattr( + MultiConnectorScheduler, + "_state_tier_sub", + lambda self: SimpleNamespace(take_state_touches_unrefreshed=lambda: 4), + ) + assert multi.take_state_touches_unrefreshed() == 4 + monkeypatch.setattr( + MultiConnectorScheduler, "_state_tier_sub", lambda self: None + ) + assert multi.take_state_touches_unrefreshed() == 0 + + def test_a_negative_round_is_refused(self): + with pytest.raises(ValueError): + StateTouchId(1, -1) diff --git a/tests/test_state_checkpoint.py b/tests/test_state_checkpoint.py index 47bc047caa..5bd859d771 100644 --- a/tests/test_state_checkpoint.py +++ b/tests/test_state_checkpoint.py @@ -19,6 +19,7 @@ import pytest from conftest import MockConfig +from atom.model_engine import state_offload as so from atom.model_engine.block_manager import AdmissionProbe, BlockManager from atom.model_engine.block_pool import BlockPool from atom.model_engine.engine_stats import EngineStats @@ -30,6 +31,7 @@ from atom.model_engine.scheduler import ScheduledBatchOutput, Scheduler from atom.model_engine.sequence import Sequence, SequenceType from atom.model_engine.state_cache import StateCache +from atom.model_engine.state_offload import StateOffloadIndex from atom.model_engine.state_pool import StateSlotPool from atom.model_engine.state_runtime import ( StateRuntime, @@ -4556,3 +4558,294 @@ def test_a_cold_start_without_a_boundary_is_untouched(self): assert seq.state_slot >= 0 assert bm.state_gate_lost_boundary == before + + +# ── state tier hit refresh: the admission hold and the voters ────────────── + + +@pytest.fixture +def clock(monkeypatch): + """`state_offload`'s monotonic clock, advanced by hand.""" + now = [1000.0] + monkeypatch.setattr(so, "monotonic", lambda: now[0]) + return now + + +#: Past every confirmation window `tier_verified_for_s` can give. +STALE = so._TIER_VERIFIED_MAX_S + 1 + + +def _paged_bm(): + bm = make_block_manager( + ckpt_config(num_kvcache_blocks=160, pool_entries={"state": 8}), + state_runtime=PAGED_COPY_RUNTIME, + ) + bm.state_offload = StateOffloadIndex(refresh_hits=True) + return bm + + +def _seq(): + return stateful_seq(list(range(40))) + + +class TestAdmissionHold: + """`record_allocation(defer_unverified=True)` holds an admission that + would load an unconfirmed image, and commits nothing while it does.""" + + def _joint(self, bm, monkeypatch, *, h=77, from_hbm=False): + from atom.model_engine.block_manager import JointBoundaryDecision + + decision = JointBoundaryDecision( + boundary_tokens=16, + boundary_hash=h, + kv_tokens=16, + claim_tokens=16, + state_from_hbm=from_hbm, + ) + commits = [] + monkeypatch.setattr(bm, "_joint_kv_boundary", lambda *a: decision) + monkeypatch.setattr(bm, "_fit_joint", lambda d, probe: d) + monkeypatch.setattr( + bm, "_commit_joint_boundary", lambda seq, d: commits.append(d) + ) + return commits + + def test_an_unconfirmed_joint_image_is_held(self, clock, monkeypatch): + bm = _paged_bm() + commits = self._joint(bm, monkeypatch) + bm.state_offload.note_stored(77) + clock[0] += STALE + seq = _seq() + + assert not bm.record_allocation( + seq, AdmissionProbe(2, [1, 2]), defer_unverified=True + ) + assert commits == [] and bm.state_offload.loads_deferred_for_verify == 1 + assert bm.state_offload.take_touches() == [77] + + bm.state_offload.apply_touch_verdict(77, True) + assert bm.record_allocation( + seq, AdmissionProbe(2, [1, 2]), defer_unverified=True + ) + assert len(commits) == 1 + + def test_a_gone_image_is_not_parked_on(self, clock, monkeypatch): + """The verdict drops the hash; the retry's boundary search no longer + sees the tier rung (`could_serve` is False), so it plans a lower one.""" + bm = _paged_bm() + self._joint(bm, monkeypatch) + bm.state_offload.note_stored(77) + clock[0] += STALE + assert not bm.record_allocation( + _seq(), AdmissionProbe(2, [1, 2]), defer_unverified=True + ) + bm.state_offload.take_touches() + bm.state_offload.apply_touch_verdict(77, False) + assert not bm._tier_can_serve(77) + + def test_a_recently_confirmed_image_goes_straight_through(self, clock, monkeypatch): + bm = _paged_bm() + commits = self._joint(bm, monkeypatch) + bm.state_offload.note_stored(77) + assert bm.record_allocation( + _seq(), AdmissionProbe(2, [1, 2]), defer_unverified=True + ) + assert len(commits) == 1 + + def test_a_state_resident_in_hbm_is_never_held(self, clock, monkeypatch): + bm = _paged_bm() + self._joint(bm, monkeypatch, from_hbm=True) + bm.state_offload.note_stored(77) + clock[0] += STALE + assert bm.record_allocation( + _seq(), AdmissionProbe(2, [1, 2]), defer_unverified=True + ) + + def test_the_fit_probe_path_never_holds(self, clock, monkeypatch): + bm = _paged_bm() + self._joint(bm, monkeypatch) + bm.state_offload.note_stored(77) + clock[0] += STALE + assert bm.record_allocation(_seq(), AdmissionProbe(2, [1, 2])) + + def test_an_hbm_hit_holds_only_for_a_tier_sourced_state(self, clock, monkeypatch): + from atom.model_engine.block_manager import JointBoundaryDecision + + bm = _paged_bm() + monkeypatch.setattr( + bm, + "_joint_kv_boundary", + lambda *a: JointBoundaryDecision(skip_reason="no_kv_reach"), + ) + monkeypatch.setattr(bm, "_commit_joint_boundary", lambda seq, d: None) + bm.state_offload.note_stored(2) + clock[0] += STALE + resident = {2} + monkeypatch.setattr( + bm.paged_state_checkpoints, "contains", lambda h: h in resident + ) + assert bm.record_allocation( + _seq(), AdmissionProbe(2, [1, 2]), defer_unverified=True + ) + resident.clear() + assert not bm.record_allocation( + _seq(), AdmissionProbe(2, [1, 2]), defer_unverified=True + ) + + def test_the_hold_judges_the_joint_fit_kept(self, clock, monkeypatch): + """A joint `_fit_joint` drops is not loaded, so its stale image must not + hold an admission that resumes at an HBM-resident hit.""" + bm = _paged_bm() + commits = self._joint(bm, monkeypatch) + monkeypatch.setattr( + bm, "_fit_joint", lambda d, probe: bm._no_joint("partial_chunk_no_room") + ) + monkeypatch.setattr(bm.paged_state_checkpoints, "contains", lambda h: h == 2) + bm.state_offload.note_stored(77) + clock[0] += STALE + assert bm.record_allocation( + _seq(), AdmissionProbe(2, [1, 2], 1), defer_unverified=True + ) + assert [d.skip_reason for d in commits] == ["partial_chunk_no_room"] + assert bm.state_offload.loads_deferred_for_verify == 0 + + +class TestVotersRegisterTouches: + """`could_serve` only answers; the voters register the hits they count.""" + + def _bm(self): + bm = _paged_bm() + bm.paged_state_checkpoints.attach_offload(bm.state_offload) + bm.state_offload.note_stored(77) + return bm + + def test_the_resumable_scan_registers_a_tier_hit(self, clock): + bm = self._bm() + assert bm.paged_state_checkpoints._reachable(77) + assert bm.state_offload.take_touches() == [77] + + def test_an_hbm_hit_is_not_a_tier_vote(self, clock, monkeypatch): + bm = self._bm() + monkeypatch.setattr( + bm.paged_state_checkpoints.store, "contains", lambda h: True + ) + assert bm.paged_state_checkpoints._reachable(77) + assert bm.state_offload.take_touches() == [] + + def test_asking_whether_the_tier_can_serve_registers_nothing(self, clock): + bm = self._bm() + assert bm._tier_can_serve(77) + assert bm.state_offload.take_touches() == [] + + def test_an_offered_load_registers_its_hash(self, clock): + bm = self._bm() + seq = _seq() + assert bm._request_state_load(seq, 77) + assert bm.state_offload.take_touches() == [77] + + +class TestResumeHash: + """`allocate` and the admission hold pick the state to resume from with + one helper, `_state_resume_hash`.""" + + def test_the_joint_boundary_wins_then_a_replay_then_the_hit(self, monkeypatch): + bm = _paged_bm() + seq = _seq() + assert bm._state_resume_hash(seq, 77, 2) == 77 + assert bm._state_resume_hash(seq, -1, 2) == 2 + assert bm._state_resume_hash(seq, -1, -1) == -1 + monkeypatch.setattr(bm, "state_replay_tokens", 8) + seq.replays_hit = True + assert bm._state_resume_hash(seq, -1, 2) == -1, "a replay restores none" + assert bm._state_resume_hash(seq, 77, 2) == 77 + + def test_a_replayed_hit_is_not_held(self, clock, monkeypatch): + from atom.model_engine.block_manager import JointBoundaryDecision + + bm = _paged_bm() + monkeypatch.setattr( + bm, + "_joint_kv_boundary", + lambda *a: JointBoundaryDecision(skip_reason="no_kv_reach"), + ) + monkeypatch.setattr(bm, "_commit_joint_boundary", lambda seq, d: None) + monkeypatch.setattr(bm.paged_state_checkpoints, "contains", lambda h: False) + monkeypatch.setattr(bm, "state_replay_tokens", 8) + bm.state_offload.note_stored(2) + clock[0] += STALE + seq = _seq() + seq.replays_hit = True + assert bm.record_allocation( + seq, AdmissionProbe(2, [1, 2]), defer_unverified=True + ) + + def test_the_hold_judges_the_hash_allocate_then_loads(self, clock, monkeypatch): + bm = _paged_bm() + offload = bm.state_offload + bm.paged_state_checkpoints.attach_offload(offload) + first = _seq() + bm.allocate(first, bm.can_allocate(first)) + bm.hash_blocks(first, bm.checkpoint_limit(first) - first.num_cached_tokens) + h = boundary_hash(bm, first) + bm.take_state_maintenance_ops() + bm.complete_previous_state_batch() + # The image left HBM; only the tier has it, and not confirmed lately. + bm.paged_state_checkpoints.unindex(h) + offload.note_stored(h) + clock[0] += STALE + + asked, attached = [], [] + verify, attach = offload.needs_verify, bm._attach_state_slots + monkeypatch.setattr( + offload, "needs_verify", lambda x: asked.append(x) or verify(x) + ) + monkeypatch.setattr( + bm, "_attach_state_slots", lambda s, x: attached.append(x) or attach(s, x) + ) + second = stateful_seq(list(range(48))) + probe = AdmissionProbe() + probe.num_cached_blocks = bm.can_allocate(second, record=False, probe=probe) + assert probe.num_cached_blocks > 0 + assert not bm.record_allocation(second, probe, defer_unverified=True) + offload.apply_touch_verdict(h, True) + assert bm.record_allocation(second, probe, defer_unverified=True) + assert bm.allocate(second, probe.num_cached_blocks) + assert set(asked) == {h} and attached == [h] + assert second.offload_joint.load_hash == h + + +def test_the_scheduler_retries_a_held_admission(monkeypatch): + """A held seq stays waiting and is admitted on a later pass.""" + sched = make_scheduler(ckpt_config()) + bm = sched.block_manager + real = bm.record_allocation + holds = [True] + + def record(seq, probe, *, defer_unverified=False): + assert defer_unverified, "admission asks to hold unverified loads" + if holds.pop() if holds else False: + return False + return real(seq, probe, defer_unverified=defer_unverified) + + monkeypatch.setattr(bm, "record_allocation", record) + seq = stateful_seq(list(range(12))) + sched.add(seq) + _, scheduled = sched.schedule() + assert scheduled == {} + assert list(sched.waiting) == [seq] and len(seq.block_table) == 0 + _, scheduled = sched.schedule() + assert seq.id in scheduled and not sched.waiting + + +def test_the_scheduler_publishes_touches_with_the_state_loads(clock): + idx = StateOffloadIndex(refresh_hits=True) + idx.note_stored(7) + idx.want_touch(7) + sent = [] + s = Scheduler.__new__(Scheduler) + s.kv_connector = SimpleNamespace( + enqueue_state_touches=lambda hs: sent.extend(hs) or True + ) + s.block_manager = SimpleNamespace(state_offload=idx) + s._publish_state_loads() + assert sent == [7] diff --git a/tests/test_state_offload_index.py b/tests/test_state_offload_index.py index e5eca4cdd9..71da0df1e8 100644 --- a/tests/test_state_offload_index.py +++ b/tests/test_state_offload_index.py @@ -1,7 +1,18 @@ # SPDX-License-Identifier: MIT -"""StateOffloadIndex must stop advertising images the worker codec dropped.""" +"""StateOffloadIndex must stop advertising images the worker codec dropped, +and keeps state images as warm in LMCache's LRU as the KV they belong to +(hit touches, their verdicts, and the admission hold).""" -from atom.model_engine.state_offload import StateOffloadIndex, state_tier_cpu_bytes +from types import SimpleNamespace + +import pytest + +from atom.model_engine import state_offload as so +from atom.model_engine.state_offload import ( + StateOffloadIndex, + state_tier_budget_warning, + state_tier_cpu_bytes, +) def _stored(idx, hashes): @@ -49,3 +60,343 @@ def test_state_tier_cpu_bytes_reads_gib(monkeypatch): assert state_tier_cpu_bytes() == 32 << 30 monkeypatch.setenv("OFFLOAD_STATE_CPU_SIZE", "96") assert state_tier_cpu_bytes() == 96 << 30 + + +@pytest.fixture +def clock(monkeypatch): + """`state_offload`'s monotonic clock, advanced by hand.""" + now = [1000.0] + monkeypatch.setattr(so, "monotonic", lambda: now[0]) + return now + + +# ── the index: votes queue touches, verdicts keep it in step ─────────────── + + +#: Past every confirmation window `tier_verified_for_s` can give. +STALE = so._TIER_VERIFIED_MAX_S + 1 + + +def _index(*hashes, refresh=True, **kw): + idx = StateOffloadIndex(refresh_hits=refresh, **kw) + for h in hashes: + idx.note_stored(h) + return idx + + +class TestIndexTouches: + def test_could_serve_queues_nothing(self, clock): + idx = _index(7) + assert idx.could_serve(7) and not idx.could_serve(8) + assert idx.take_touches() == [] + + def test_a_registered_hit_queues_a_touch(self, clock): + idx = _index(7) + idx.want_touch(7) + idx.want_touch(8) # not indexed: nothing to touch + assert idx.take_touches() == [7] + assert idx.take_touches() == [] + + def test_no_touches_without_refresh(self, clock): + idx = _index(7, refresh=False) + assert idx.could_serve(7) + idx.want_touch(7) + assert idx.take_touches() == [] + + def test_a_hash_voted_every_pass_is_touched_at_an_interval(self, clock): + idx = _index(7) + idx.want_touch(7) + assert idx.take_touches() == [7] + clock[0] += so._TIER_TOUCH_INTERVAL_S / 2 + idx.want_touch(7) + assert idx.take_touches() == [] + clock[0] += so._TIER_TOUCH_INTERVAL_S + idx.want_touch(7) + assert idx.take_touches() == [7] + assert idx.touches_sent == 2 + + def test_a_gone_verdict_drops_the_hash(self, clock): + idx = _index(7) + clock[0] += 1 + idx.want_touch(7) + assert idx.take_touches() == [7] + idx.apply_touch_verdict(7, False) + assert not idx.could_serve(7) + assert idx.touch_misses == 1 + assert idx.stats()["touch_misses"] == 1 + + def test_a_gone_verdict_older_than_a_store_is_ignored(self, clock): + """The touch looked before the store's bytes landed; its verdict must + not erase the image that store just indexed.""" + idx = _index(7) + clock[0] += 1 + idx.want_touch(7) + assert idx.take_touches() == [7] + clock[0] += 0.01 + idx.note_stored(7) + idx.apply_touch_verdict(7, False) + assert idx.could_serve(7) + assert idx.touch_misses == 0 + + def test_a_gone_verdict_for_a_forgotten_then_restored_hash_is_ignored(self, clock): + idx = _index(7) + clock[0] += 1 + idx.want_touch(7) + assert idx.take_touches() == [7] + idx.forget(7) + clock[0] += 0.01 + idx.note_stored(7) + idx.apply_touch_verdict(7, False) + assert idx.could_serve(7) + + def test_a_gone_verdict_after_a_store_and_a_newer_touch_applies(self, clock): + idx = _index(7) + clock[0] += 1 + idx.note_stored(7) + clock[0] += 1 + idx.want_touch(7) + assert idx.take_touches() == [7] + idx.apply_touch_verdict(7, False) + assert not idx.could_serve(7) + + def test_a_resident_verdict_refreshes_the_cap_order(self, clock): + idx = StateOffloadIndex(max_cpu_entries=19, refresh_hits=True) # cap 3 + for h in (1, 2, 3): + idx.note_stored(h) + idx.apply_touch_verdict(1, True) + idx.note_stored(4) + assert idx.could_serve(1) and not idx.could_serve(2) + + def test_a_verdict_for_a_forgotten_hash_is_ignored(self, clock): + idx = _index(7) + idx.forget(7) + idx.apply_touch_verdict(7, True) + assert not idx.could_serve(7) + + def test_stop_refreshing_drops_pending_touches(self, clock): + idx = _index(7) + idx.want_touch(7) + idx.stop_refreshing() + assert idx.take_touches() == [] + assert not idx.needs_verify(7) + + +class TestNeedsVerify: + def test_a_fresh_store_needs_no_verdict(self, clock): + assert not _index(7).needs_verify(7) + + def test_a_store_only_role_holds_nothing(self, clock): + """A `kv_producer` indexes its own stores but never loads them.""" + idx = _index(7, can_load=False) + clock[0] += STALE + assert not idx.needs_verify(7) + assert idx.take_touches() == [] + + def test_a_stale_hash_waits_for_its_touch(self, clock): + idx = _index(7) + clock[0] += STALE + assert idx.needs_verify(7) + assert idx.take_touches() == [7], "the hold asks for its own touch" + assert idx.needs_verify(7), "still out" + idx.apply_touch_verdict(7, True) + assert not idx.needs_verify(7) + + def test_a_gone_verdict_ends_the_hold_by_dropping_the_hash(self, clock): + idx = _index(7) + clock[0] += STALE + assert idx.needs_verify(7) + idx.take_touches() + idx.apply_touch_verdict(7, False) + assert not idx.needs_verify(7) and not idx.could_serve(7) + + def test_a_lost_verdict_holds_no_longer_than_the_lookup_defer(self, clock): + idx = _index(7, defer_s=2) + clock[0] += STALE + # A held request re-probes every pass and its touch is re-sent every + # interval; verdicts never come back. Re-sends must not extend the hold. + held_for, step = 0.0, 0.05 + while idx.needs_verify(7): + assert held_for <= 2 + 1e-6, "the hold outlived OFFLOAD_LOOKUP_DEFER_S" + assert idx.could_serve(7) + idx.take_touches() + clock[0] += step + held_for += step + assert held_for >= 2 + assert idx.touches_sent > 1, "the touch was re-sent while held" + assert idx.verify_holds_expired == 1 + assert idx.stats()["verify_holds_expired"] == 1 + assert not idx.needs_verify(7), "later admissions are not held again" + assert idx.verify_holds_expired == 1 + + def test_evidence_after_an_expired_hold_rearms_it(self, clock): + idx = _index(7, defer_s=2) + clock[0] += STALE + assert idx.needs_verify(7) + clock[0] += 3 + assert not idx.needs_verify(7) + idx.apply_touch_verdict(7, True) + clock[0] += STALE + assert idx.needs_verify(7), "a fresh hold, timed from now" + clock[0] += 1 + assert idx.needs_verify(7) + + def test_a_verdict_ends_the_hold_timer(self, clock): + idx = _index(7, defer_s=2) + clock[0] += STALE + assert idx.needs_verify(7) + clock[0] += 1.5 + idx.apply_touch_verdict(7, True) + clock[0] += STALE + assert idx.needs_verify(7) + clock[0] += 1.5 + assert idx.needs_verify(7), "the earlier 1.5 s does not count" + assert idx.verify_holds_expired == 0 + + def test_a_zero_defer_never_holds(self, clock, monkeypatch): + monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "0") + idx = _index(7) + assert idx.defer_s == 0 + clock[0] += STALE + assert not idx.needs_verify(7) + + def test_the_lookup_defer_is_read_once_at_construction(self, clock, monkeypatch): + """`needs_verify` is on the admission path; it must not re-read and + re-parse `OFFLOAD_LOOKUP_DEFER_S` on every call.""" + monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "2") + idx = _index(7) + assert idx.defer_s == 2.0 + + def unread(): + raise AssertionError("OFFLOAD_LOOKUP_DEFER_S read on the admission path") + + monkeypatch.setattr(so, "verify_defer_s", unread) + monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "0") + clock[0] += STALE + assert idx.needs_verify(7), "the value read at construction still holds" + + def test_an_index_without_refresh_never_reads_the_defer(self, monkeypatch): + def unread(): + raise AssertionError("read for an index that never holds") + + monkeypatch.setattr(so, "verify_defer_s", unread) + assert StateOffloadIndex().defer_s == 0 + + def test_a_completed_load_counts_as_verified(self, clock): + idx = _index(7) + clock[0] += STALE + assert idx.request_load("r", 7) + idx.complete_load("r") + assert not idx.needs_verify(7) + + +class TestVerifiedWindow: + """How long a confirmed image is trusted scales with the LMCache pool it + has to age out of, within [5 s, 60 s].""" + + GIB = 1 << 30 + + def test_the_agentx_pool_keeps_about_the_measured_window(self): + assert so.tier_verified_for_s(int(208.6 * self.GIB)) == pytest.approx(52.15) + + def test_the_window_grows_with_the_pool_between_its_bounds(self): + sizes = [8, 20, 32, 64, 100, 208.6, 240, 512, 2048] + windows = [so.tier_verified_for_s(int(g * self.GIB)) for g in sizes] + assert windows == sorted(windows) + assert windows[0] == windows[1] == so._TIER_VERIFIED_MIN_S + assert windows[2] == pytest.approx(8.0) + assert windows[-3] == windows[-1] == so._TIER_VERIFIED_MAX_S + + def test_an_unknown_pool_gets_the_floor(self): + assert so.tier_verified_for_s(0) == so._TIER_VERIFIED_MIN_S + + def test_the_index_holds_by_its_own_window(self, clock): + idx = _index(7, verified_for_s=so.tier_verified_for_s(32 * self.GIB), defer_s=2) + assert idx.verified_for_s == pytest.approx(8.0) + clock[0] += 7.5 + assert not idx.needs_verify(7) + clock[0] += 1 + assert idx.needs_verify(7) + + +class TestTouchPlumbing: + """The index's two calls on the connector, made by the scheduler once a + pass (`publish_touches` before the metadata is built, + `collect_touch_verdicts` before the next admission).""" + + def test_voted_hashes_reach_the_connector(self, clock): + idx = _index(7) + idx.want_touch(7) + sent = [] + idx.publish_touches( + SimpleNamespace(enqueue_state_touches=lambda hs: sent.extend(hs) or True) + ) + assert sent == [7] and idx.refresh_hits + + def test_a_connector_without_touches_turns_the_refresh_off(self, clock): + idx = _index(7) + idx.want_touch(7) + idx.publish_touches(SimpleNamespace()) + assert not idx.refresh_hits + clock[0] += STALE + assert not idx.needs_verify(7), "nothing would ever answer the hold" + + def test_verdicts_are_applied_to_the_index(self, clock): + idx = _index(7, 8) + clock[0] += 1 + idx.want_touch(7) + idx.want_touch(8) + assert sorted(idx.take_touches()) == [7, 8] + idx.collect_touch_verdicts( + SimpleNamespace(take_state_touch_verdicts=lambda: {7: False, 8: True}) + ) + assert not idx.could_serve(7) and idx.could_serve(8) + + def test_presence_only_touches_reach_the_stats_line(self, clock): + idx = _index(7) + connector = SimpleNamespace( + take_state_touch_verdicts=lambda: {7: True}, + take_state_touches_unrefreshed=lambda: 3, + ) + assert idx.stats()["touches_unrefreshed"] == 0 + idx.collect_touch_verdicts(connector) + idx.collect_touch_verdicts(connector) + assert idx.stats()["touches_unrefreshed"] == 6 + + def test_an_offered_load_registers_a_touch(self, clock): + idx = _index(7) + assert idx.request_load("r", 7) + assert idx.take_touches() == [7] + + def test_a_hold_is_counted(self, clock): + idx = _index(7) + assert not idx.hold_admission(-1), "nothing to load" + assert not idx.hold_admission(7), "confirmed by its store" + clock[0] += STALE + assert idx.hold_admission(7) + assert idx.loads_deferred_for_verify == 1 + assert idx.stats()["loads_deferred_for_verify"] == 1 + + +# ── the startup budget check ─────────────────────────────────────────────── + + +GIB = 1 << 30 + + +class TestBudgetWarning: + def test_a_cap_below_half_the_pool_is_fine(self): + assert state_tier_budget_warning(59 * GIB, 208 * GIB, 112 << 20) is None + + def test_a_cap_the_size_of_the_pool_bounds_nothing(self): + """#2562's accuracy run: a 32 GiB pool with the 32 GiB default cap.""" + msg = state_tier_budget_warning(32 * GIB, 32 * GIB, 112 << 20) + assert "bounds nothing" in msg + assert "LMCACHE_MAX_LOCAL_CPU_SIZE=32.0" in msg + + def test_a_cap_over_half_the_pool_is_called_out(self): + msg = state_tier_budget_warning(40 * GIB, 64 * GIB, 112 << 20) + assert msg is not None and "more than half" in msg + + def test_unknown_sizes_are_not_judged(self): + assert state_tier_budget_warning(32 * GIB, 0, 1) is None + assert state_tier_budget_warning(0, 32 * GIB, 1) is None diff --git a/tests/test_state_tier_lru_refresh.py b/tests/test_state_tier_lru_refresh.py deleted file mode 100644 index 2aded0a403..0000000000 --- a/tests/test_state_tier_lru_refresh.py +++ /dev/null @@ -1,1224 +0,0 @@ -# SPDX-License-Identifier: MIT -"""State images are refreshed in LMCache's LRU like the KV they belong to. - -State images and paged KV share one LocalCPU pool. A KV lookup refreshes every -chunk it touches (`contains(pin=True)` then `touch_cache()` -> -`cache_policy.update_on_hit`), but a state hit is answered engine-side by -`StateOffloadIndex`, so LMCache never learned an image was used: a prefix hit -on every turn kept its KV at the young end and its image at the cold end, where -pool pressure evicted it first, and the index kept advertising it until a joint -load parked on it, missed, and recomputed the whole prompt. - -These tests model LocalCPU's hot cache and LRU policy, and cover the three -pieces of the fix: the codec touch, the engine-side touch/verdict loop that -keeps the index in step with evictions, and the admission hold that lets a -verdict land before a request parks on an image. -""" - -import dataclasses -import logging -import sys -import threading -import types -from collections import OrderedDict -from types import SimpleNamespace - -import pytest - -from atom.kv_transfer.disaggregation.types import ConnectorCompletion, StateTouchId -from atom.kv_transfer.offload.hybrid.kimi_k3.state_object import TouchResult -from atom.model_engine import state_offload as so -from atom.model_engine.block_manager import AdmissionProbe -from atom.model_engine.state_offload import ( - StateOffloadIndex, - state_tier_budget_warning, -) - - -@pytest.fixture -def lmcache_key(monkeypatch): - @dataclasses.dataclass(frozen=True) - class CacheEngineKey: - model_name: str - world_size: int - worker_id: int - chunk_hash: int - dtype: object - - utils_module = types.ModuleType("lmcache.utils") - utils_module.CacheEngineKey = CacheEngineKey - lmcache_module = types.ModuleType("lmcache") - lmcache_module.__path__ = [] - lmcache_module.utils = utils_module - monkeypatch.setitem(sys.modules, "lmcache", lmcache_module) - monkeypatch.setitem(sys.modules, "lmcache.utils", utils_module) - return CacheEngineKey - - -@pytest.fixture -def clock(monkeypatch): - """`state_offload`'s monotonic clock, advanced by hand.""" - now = [1000.0] - monkeypatch.setattr(so, "monotonic", lambda: now[0]) - return now - - -class _Obj: - """A MemoryObj's reference count, and nothing else.""" - - def __init__(self, size: int = 1024) -> None: - self.refs = 1 - self.size = size - - def ref_count_up(self) -> None: - self.refs += 1 - - def ref_count_down(self) -> None: - assert self.refs > 0, "reference discharged twice" - self.refs -= 1 - - def get_ref_count(self) -> int: - return self.refs - - def get_size(self) -> int: - return self.size - - -class _LRUPolicy: - """LMCache's `LRUCachePolicy.update_on_hit`.""" - - def update_on_hit(self, key, cache) -> None: - cache.move_to_end(key) - - -class _LocalCPU: - """`storage_manager` over LocalCPU's rules. - - `batched_put` adopts an object with its own reference unless the key is - already resident (LocalCPU refuses duplicates), then drops the caller's. - `get` does not touch the LRU (`get_blocking` only takes a reference). - `evict_oldest` is the pool making room for KV: oldest first, skipping - anything another holder still references. - """ - - def __init__(self, refuse_all: bool = False) -> None: - self.hot: OrderedDict = OrderedDict() - # A second backend (LocalDisk) for presence only: `contains` answers - # from it, `get` does not read it (no test loads a disk-only image). - self.disk: dict = {} - self.refuse_all = refuse_all - self.fail_put = False - self.puts = 0 - self.local_cpu_backend = SimpleNamespace( - cpu_lock=threading.Lock(), hot_cache=self.hot, cache_policy=_LRUPolicy() - ) - - def batched_put(self, keys, objs) -> None: - if self.fail_put: - raise RuntimeError("put failed") - self.puts += 1 - for k, o in zip(keys, objs, strict=True): - if not self.refuse_all and k not in self.hot: - o.ref_count_up() - self.hot[k] = o - o.ref_count_down() - - def get(self, key): - o = self.hot.get(key) - if o is not None: - o.ref_count_up() - return o - - def contains(self, key): - if key in self.hot: - return "LocalCPUBackend" - return "LocalDiskBackend" if key in self.disk else None - - def spill_to_disk(self, key) -> None: - """LocalCPU evicted `key`; the disk backend still holds a copy.""" - self.disk[key] = True - self.hot.pop(key).ref_count_down() - - def remove(self, key) -> int: - o = self.hot.pop(key, None) - if o is None: - return 0 - o.ref_count_down() - return 1 - - def evict_oldest(self): - for k, o in self.hot.items(): - if o.get_ref_count() == 1: - self.hot.pop(k).ref_count_down() - return k - return None - - -def _codec(storage, max_entries): - from atom.kv_transfer.offload.hybrid.kimi_k3.state_object import StateByteCodec - - codec = StateByteCodec( - backend=SimpleNamespace( - page_unit_views=lambda units: [], state_entry_views=lambda slot: [] - ), - staged=SimpleNamespace(pack=lambda views, obj: None, unpack=lambda o, v: None), - entry_bytes=1024, - model_name="qwen", - world_size=1, - worker_id=0, - layout_id="L", - max_cache_bytes=None if max_entries is None else max_entries * 1024, - ) - codec.bind_storage_manager(storage) - allocated = [] - - def allocate(nbytes): - allocated.append(_Obj(nbytes)) - return allocated[-1] - - codec._allocate = allocate - return codec, allocated - - -def _kv(pool, name): - pool.hot[name] = _Obj() - - -# ── the codec: touch, get and re-store refresh LMCache's LRU ─────────────── - - -class TestCodecRefresh: - def test_an_untouched_image_is_the_first_thing_pressure_evicts(self, lmcache_key): - """The bug: KV written after the image is younger, however often the - prefix hits.""" - pool = _LocalCPU() - codec, _ = _codec(pool, max_entries=None) - assert codec.put(1, []) - _kv(pool, "kv-chunk") - assert pool.evict_oldest() == codec.key(1) - - def test_a_touch_moves_the_image_past_older_kv(self, lmcache_key): - pool = _LocalCPU() - codec, _ = _codec(pool, max_entries=None) - assert codec.put(1, []) - _kv(pool, "kv-chunk") - assert codec.touch(1) is TouchResult.REFRESHED - assert pool.evict_oldest() == "kv-chunk" - assert codec.key(1) in pool.hot - - def test_a_touch_reports_an_evicted_image_gone(self, lmcache_key): - pool = _LocalCPU() - codec, _ = _codec(pool, max_entries=2) - assert codec.put(1, []) - assert pool.evict_oldest() == codec.key(1) - assert codec.touch(1) is TouchResult.GONE - assert codec.key(1) not in codec._cache_keys, "no longer counts against the cap" - - def test_a_resume_refreshes_the_image(self, lmcache_key): - """`get_blocking` takes a reference and nothing else.""" - pool = _LocalCPU() - codec, objs = _codec(pool, max_entries=None) - assert codec.put(1, []) - _kv(pool, "kv-chunk") - assert codec.get(1, 0) - assert objs[0].refs == 1, "the load's reference is discharged" - assert pool.evict_oldest() == "kv-chunk" - - def test_a_missed_get_leaves_the_cap_lru(self, lmcache_key): - pool = _LocalCPU() - codec, _ = _codec(pool, max_entries=2) - assert codec.put(1, []) - pool.evict_oldest() - assert not codec.get(1, 0) - assert codec._cache_keys == OrderedDict() - - def test_a_resident_image_is_refreshed_not_stored_twice(self, lmcache_key): - pool = _LocalCPU() - codec, objs = _codec(pool, max_entries=2) - released = [] - assert codec.put(1, []) - _kv(pool, "kv-chunk") - assert codec.put(1, [], on_source_released=lambda: released.append(True)) - - assert len(objs) == 1, "no second image-sized allocation and D2H" - assert released == [True], "the source units must still be handed back" - assert pool.evict_oldest() == "kv-chunk" - - def test_an_evicted_image_is_stored_again(self, lmcache_key): - pool = _LocalCPU() - codec, objs = _codec(pool, max_entries=2) - assert codec.put(1, []) - pool.evict_oldest() - assert codec.put(1, []) - assert len(objs) == 2 and codec.key(1) in pool.hot - - def test_the_cap_retires_the_least_recently_used_image(self, lmcache_key): - """Not the oldest store: the image a touch kept warm survives.""" - pool = _LocalCPU() - codec, objs = _codec(pool, max_entries=2) - assert codec.put(1, []) and codec.put(2, []) - assert codec.touch(1) is TouchResult.REFRESHED - assert codec.put(3, []) - - assert codec.key(2) not in pool.hot - assert objs[1].refs == 0, "a retired image goes back to the allocator" - assert list(codec._cache_keys) == [codec.key(1), codec.key(3)] - - def test_images_hold_no_extra_reference(self, lmcache_key): - """No reservation: KV pressure may evict an image like any chunk.""" - pool = _LocalCPU() - codec, objs = _codec(pool, max_entries=2) - assert codec.put(1, []) - assert objs[0].refs == 1 - assert pool.evict_oldest() == codec.key(1) - - def test_a_refused_put_keeps_no_reference(self, lmcache_key): - pool = _LocalCPU(refuse_all=True) - codec, objs = _codec(pool, max_entries=2) - assert codec.put(1, []) - assert objs[0].refs == 0, "holding a refused object would strand its block" - - def test_a_failed_put_discharges_its_reference(self, lmcache_key): - pool = _LocalCPU() - pool.fail_put = True - codec, objs = _codec(pool, max_entries=2) - with pytest.raises(RuntimeError): - codec.put(1, []) - assert objs[0].refs == 0 - - def test_without_a_local_cpu_backend_a_touch_asks_presence( - self, lmcache_key, caplog - ): - pool = _LocalCPU() - pool.local_cpu_backend = None - with caplog.at_level(logging.WARNING, logger="atom"): - codec, _ = _codec(pool, max_entries=None) - assert any("no local_cpu_backend" in r.message for r in caplog.records) - pool.hot[codec.key(1)] = _Obj() - assert codec.touch(1) is TouchResult.PRESENT - pool.hot.clear() - assert codec.touch(1) is TouchResult.GONE - - def _assert_degrades(self, pool, missing, caplog): - """Bind warns once naming what is missing; touches then only confirm - presence, and each one is counted rather than passing as a refresh.""" - with caplog.at_level(logging.WARNING, logger="atom"): - codec, _ = _codec(pool, max_entries=None) - bind = [r.message for r in caplog.records if missing in r.message] - assert len(bind) == 1 and "state_offload_touches_unrefreshed" in bind[0] - assert codec.put(1, []) - assert codec.touch(1) is TouchResult.PRESENT - assert codec.touch(1) is TouchResult.PRESENT - assert codec._unrefreshed_touches == 2 - - @pytest.mark.parametrize("dropped", ["cpu_lock", "hot_cache", "cache_policy"]) - def test_a_local_cpu_without_the_lru_internals_degrades_to_presence( - self, lmcache_key, dropped, caplog - ): - """An LMCache whose internals moved keeps serving loads; the lost - refresh shows in the stats instead of taking the engine down.""" - pool = _LocalCPU() - setattr(pool.local_cpu_backend, dropped, None) - self._assert_degrades(pool, f"local_cpu_backend.{dropped}", caplog) - - def test_a_policy_without_update_on_hit_degrades_to_presence( - self, lmcache_key, caplog - ): - pool = _LocalCPU() - pool.local_cpu_backend.cache_policy = SimpleNamespace() - self._assert_degrades(pool, "cache_policy.update_on_hit", caplog) - - def test_a_storage_manager_without_local_cpu_backend_degrades_to_presence( - self, lmcache_key, caplog - ): - pool = _LocalCPU() - del pool.local_cpu_backend - self._assert_degrades(pool, "no local_cpu_backend", caplog) - - def test_the_lru_is_resolved_once_at_bind(self, lmcache_key): - pool = _LocalCPU() - codec, _ = _codec(pool, max_entries=None) - assert codec.put(1, []) - _kv(pool, "kv-chunk") - pool.local_cpu_backend = None - assert codec.touch(1) is TouchResult.REFRESHED - assert pool.evict_oldest() == "kv-chunk" - - def test_a_disk_only_image_is_counted_and_warned_once(self, lmcache_key, caplog): - """Present, so still readable, but its LRU was not refreshed: it must - not be silent.""" - pool = _LocalCPU() - codec, _ = _codec(pool, max_entries=None) - assert codec.put(1, []) - pool.spill_to_disk(codec.key(1)) - with caplog.at_level(logging.WARNING, logger="atom"): - assert codec.touch(1) is TouchResult.PRESENT - assert codec.touch(1) is TouchResult.PRESENT - assert codec._unrefreshed_touches == 2 - warned = [r for r in caplog.records if "touches_unrefreshed" in r.message] - assert len(warned) == 1 - - def test_a_refreshed_or_gone_image_is_not_counted(self, lmcache_key, caplog): - pool = _LocalCPU() - codec, _ = _codec(pool, max_entries=None) - assert codec.put(1, []) - with caplog.at_level(logging.WARNING, logger="atom"): - assert codec.touch(1) is TouchResult.REFRESHED - pool.hot.clear() - assert codec.touch(1) is TouchResult.GONE - assert codec._unrefreshed_touches == 0 - assert not [r for r in caplog.records if "touches_unrefreshed" in r.message] - - def test_a_disk_only_image_is_stored_back_into_the_cpu_tier(self, lmcache_key): - """The put fast path is for an image LocalCPU still holds, not for - one only the disk does.""" - pool = _LocalCPU() - codec, objs = _codec(pool, max_entries=None) - assert codec.put(1, []) - pool.spill_to_disk(codec.key(1)) - released = [] - assert codec.put(1, [], on_source_released=lambda: released.append(True)) - assert pool.puts == 2 and len(objs) == 2 - assert codec.key(1) in pool.hot - assert released == [True] - assert codec._unrefreshed_touches == 0, "a store is not a touch verdict" - - -# ── the index: votes queue touches, verdicts keep it in step ─────────────── - - -#: Past every confirmation window `tier_verified_for_s` can give. -STALE = so._TIER_VERIFIED_MAX_S + 1 - - -def _index(*hashes, refresh=True, **kw): - idx = StateOffloadIndex(refresh_hits=refresh, **kw) - for h in hashes: - idx.note_stored(h) - return idx - - -class TestIndexTouches: - def test_could_serve_queues_nothing(self, clock): - idx = _index(7) - assert idx.could_serve(7) and not idx.could_serve(8) - assert idx.take_touches() == [] - - def test_a_registered_hit_queues_a_touch(self, clock): - idx = _index(7) - idx.want_touch(7) - idx.want_touch(8) # not indexed: nothing to touch - assert idx.take_touches() == [7] - assert idx.take_touches() == [] - - def test_no_touches_without_refresh(self, clock): - idx = _index(7, refresh=False) - assert idx.could_serve(7) - idx.want_touch(7) - assert idx.take_touches() == [] - - def test_a_hash_voted_every_pass_is_touched_at_an_interval(self, clock): - idx = _index(7) - idx.want_touch(7) - assert idx.take_touches() == [7] - clock[0] += so._TIER_TOUCH_INTERVAL_S / 2 - idx.want_touch(7) - assert idx.take_touches() == [] - clock[0] += so._TIER_TOUCH_INTERVAL_S - idx.want_touch(7) - assert idx.take_touches() == [7] - assert idx.touches_sent == 2 - - def test_a_gone_verdict_drops_the_hash(self, clock): - idx = _index(7) - clock[0] += 1 - idx.want_touch(7) - assert idx.take_touches() == [7] - idx.apply_touch_verdict(7, False) - assert not idx.could_serve(7) - assert idx.touch_misses == 1 - assert idx.stats()["touch_misses"] == 1 - - def test_a_gone_verdict_older_than_a_store_is_ignored(self, clock): - """The touch looked before the store's bytes landed; its verdict must - not erase the image that store just indexed.""" - idx = _index(7) - clock[0] += 1 - idx.want_touch(7) - assert idx.take_touches() == [7] - clock[0] += 0.01 - idx.note_stored(7) - idx.apply_touch_verdict(7, False) - assert idx.could_serve(7) - assert idx.touch_misses == 0 - - def test_a_gone_verdict_for_a_forgotten_then_restored_hash_is_ignored(self, clock): - idx = _index(7) - clock[0] += 1 - idx.want_touch(7) - assert idx.take_touches() == [7] - idx.forget(7) - clock[0] += 0.01 - idx.note_stored(7) - idx.apply_touch_verdict(7, False) - assert idx.could_serve(7) - - def test_a_gone_verdict_after_a_store_and_a_newer_touch_applies(self, clock): - idx = _index(7) - clock[0] += 1 - idx.note_stored(7) - clock[0] += 1 - idx.want_touch(7) - assert idx.take_touches() == [7] - idx.apply_touch_verdict(7, False) - assert not idx.could_serve(7) - - def test_a_resident_verdict_refreshes_the_cap_order(self, clock): - idx = StateOffloadIndex(max_cpu_entries=19, refresh_hits=True) # cap 3 - for h in (1, 2, 3): - idx.note_stored(h) - idx.apply_touch_verdict(1, True) - idx.note_stored(4) - assert idx.could_serve(1) and not idx.could_serve(2) - - def test_a_verdict_for_a_forgotten_hash_is_ignored(self, clock): - idx = _index(7) - idx.forget(7) - idx.apply_touch_verdict(7, True) - assert not idx.could_serve(7) - - def test_stop_refreshing_drops_pending_touches(self, clock): - idx = _index(7) - idx.want_touch(7) - idx.stop_refreshing() - assert idx.take_touches() == [] - assert not idx.needs_verify(7) - - -class TestNeedsVerify: - def test_a_fresh_store_needs_no_verdict(self, clock): - assert not _index(7).needs_verify(7) - - def test_a_store_only_role_holds_nothing(self, clock): - """A `kv_producer` indexes its own stores but never loads them.""" - idx = _index(7, can_load=False) - clock[0] += STALE - assert not idx.needs_verify(7) - assert idx.take_touches() == [] - - def test_a_stale_hash_waits_for_its_touch(self, clock): - idx = _index(7) - clock[0] += STALE - assert idx.needs_verify(7) - assert idx.take_touches() == [7], "the hold asks for its own touch" - assert idx.needs_verify(7), "still out" - idx.apply_touch_verdict(7, True) - assert not idx.needs_verify(7) - - def test_a_gone_verdict_ends_the_hold_by_dropping_the_hash(self, clock): - idx = _index(7) - clock[0] += STALE - assert idx.needs_verify(7) - idx.take_touches() - idx.apply_touch_verdict(7, False) - assert not idx.needs_verify(7) and not idx.could_serve(7) - - def test_a_lost_verdict_holds_no_longer_than_the_lookup_defer(self, clock): - idx = _index(7, defer_s=2) - clock[0] += STALE - # A held request re-probes every pass and its touch is re-sent every - # interval; verdicts never come back. Re-sends must not extend the hold. - held_for, step = 0.0, 0.05 - while idx.needs_verify(7): - assert held_for <= 2 + 1e-6, "the hold outlived OFFLOAD_LOOKUP_DEFER_S" - assert idx.could_serve(7) - idx.take_touches() - clock[0] += step - held_for += step - assert held_for >= 2 - assert idx.touches_sent > 1, "the touch was re-sent while held" - assert idx.verify_holds_expired == 1 - assert idx.stats()["verify_holds_expired"] == 1 - assert not idx.needs_verify(7), "later admissions are not held again" - assert idx.verify_holds_expired == 1 - - def test_evidence_after_an_expired_hold_rearms_it(self, clock): - idx = _index(7, defer_s=2) - clock[0] += STALE - assert idx.needs_verify(7) - clock[0] += 3 - assert not idx.needs_verify(7) - idx.apply_touch_verdict(7, True) - clock[0] += STALE - assert idx.needs_verify(7), "a fresh hold, timed from now" - clock[0] += 1 - assert idx.needs_verify(7) - - def test_a_verdict_ends_the_hold_timer(self, clock): - idx = _index(7, defer_s=2) - clock[0] += STALE - assert idx.needs_verify(7) - clock[0] += 1.5 - idx.apply_touch_verdict(7, True) - clock[0] += STALE - assert idx.needs_verify(7) - clock[0] += 1.5 - assert idx.needs_verify(7), "the earlier 1.5 s does not count" - assert idx.verify_holds_expired == 0 - - def test_a_zero_defer_never_holds(self, clock, monkeypatch): - monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "0") - idx = _index(7) - assert idx.defer_s == 0 - clock[0] += STALE - assert not idx.needs_verify(7) - - def test_the_lookup_defer_is_read_once_at_construction(self, clock, monkeypatch): - """`needs_verify` is on the admission path; it must not re-read and - re-parse `OFFLOAD_LOOKUP_DEFER_S` on every call.""" - monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "2") - idx = _index(7) - assert idx.defer_s == 2.0 - - def unread(): - raise AssertionError("OFFLOAD_LOOKUP_DEFER_S read on the admission path") - - monkeypatch.setattr(so, "verify_defer_s", unread) - monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "0") - clock[0] += STALE - assert idx.needs_verify(7), "the value read at construction still holds" - - def test_an_index_without_refresh_never_reads_the_defer(self, monkeypatch): - def unread(): - raise AssertionError("read for an index that never holds") - - monkeypatch.setattr(so, "verify_defer_s", unread) - assert StateOffloadIndex().defer_s == 0 - - def test_a_completed_load_counts_as_verified(self, clock): - idx = _index(7) - clock[0] += STALE - assert idx.request_load("r", 7) - idx.complete_load("r") - assert not idx.needs_verify(7) - - -class TestVerifiedWindow: - """How long a confirmed image is trusted scales with the LMCache pool it - has to age out of, within [5 s, 60 s].""" - - GIB = 1 << 30 - - def test_the_agentx_pool_keeps_about_the_measured_window(self): - assert so.tier_verified_for_s(int(208.6 * self.GIB)) == pytest.approx(52.15) - - def test_the_window_grows_with_the_pool_between_its_bounds(self): - sizes = [8, 20, 32, 64, 100, 208.6, 240, 512, 2048] - windows = [so.tier_verified_for_s(int(g * self.GIB)) for g in sizes] - assert windows == sorted(windows) - assert windows[0] == windows[1] == so._TIER_VERIFIED_MIN_S - assert windows[2] == pytest.approx(8.0) - assert windows[-3] == windows[-1] == so._TIER_VERIFIED_MAX_S - - def test_an_unknown_pool_gets_the_floor(self): - assert so.tier_verified_for_s(0) == so._TIER_VERIFIED_MIN_S - - def test_the_index_holds_by_its_own_window(self, clock): - idx = _index(7, verified_for_s=so.tier_verified_for_s(32 * self.GIB), defer_s=2) - assert idx.verified_for_s == pytest.approx(8.0) - clock[0] += 7.5 - assert not idx.needs_verify(7) - clock[0] += 1 - assert idx.needs_verify(7) - - def test_the_block_manager_sizes_the_window_from_the_lmcache_pool( - self, monkeypatch - ): - from test_block_manager import _OFFLOAD_KVC, _bm_with_state_tier - - monkeypatch.setattr( - so, - "_lmcache_engine_config", - lambda config: SimpleNamespace(chunk_size=256, max_local_cpu_size=100), - ) - bm = _bm_with_state_tier(monkeypatch, _OFFLOAD_KVC) - assert bm._joint_chunk_tokens == 256 - assert bm.state_offload.verified_for_s == pytest.approx(25.0) - - def test_an_unreadable_lmcache_config_gets_the_floor(self, monkeypatch, caplog): - import sys - - from test_block_manager import _OFFLOAD_KVC, _bm_with_state_tier - - monkeypatch.setitem(sys.modules, "lmcache.v1.config", None) - with caplog.at_level(logging.WARNING, logger="atom"): - bm = _bm_with_state_tier(monkeypatch, _OFFLOAD_KVC) - assert bm._joint_chunk_tokens == 0 - assert bm.state_offload.verified_for_s == so._TIER_VERIFIED_MIN_S - assert sum(r.exc_info is not None for r in caplog.records) == 1 - - -# ── the connector: touches out in the metadata, verdicts back by quorum ──── - - -def _k3_worker(tier): - from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( - KimiK3OffloadConnector, - ) - - c = KimiK3OffloadConnector.__new__(KimiK3OffloadConnector) - c._state_tier = tier - c._touch_verdicts = {} - c._touch_failures = 0 - return c - - -class TestConnectorTouches: - def test_the_worker_touches_and_records_each_verdict(self): - touched = [] - - found = {1: TouchResult.REFRESHED, 2: TouchResult.GONE, 3: TouchResult.PRESENT} - - def touch(h): - touched.append(h) - return found[h] - - worker = _k3_worker(SimpleNamespace(codec=SimpleNamespace(touch=touch))) - ids = [StateTouchId(1, 0), StateTouchId(2, 1), StateTouchId(3, 2)] - worker._touch_states(SimpleNamespace(state_touches=ids)) - assert touched == [1, 2, 3] - # (resident, refreshed): only a present-but-unrefreshed image is False - # on the second. - assert worker._touch_verdicts == { - ids[0]: (True, True), - ids[1]: (False, True), - ids[2]: (True, False), - } - - def test_with_no_tier_every_image_is_gone(self): - worker = _k3_worker(None) - worker._touch_states(SimpleNamespace(state_touches=[StateTouchId(1, 0)])) - assert worker._touch_verdicts == {StateTouchId(1, 0): (False, True)} - - def test_a_failing_touch_does_not_erase_the_index(self): - def touch(h): - raise RuntimeError("backend hiccup") - - worker = _k3_worker(SimpleNamespace(codec=SimpleNamespace(touch=touch))) - worker._touch_states(SimpleNamespace(state_touches=[StateTouchId(1, 0)])) - assert worker._touch_verdicts[StateTouchId(1, 0)][0] is True - - def test_a_failing_touch_counts_unrefreshed_and_logs_once(self, caplog): - """A drifted LMCache raises on every touch: that must show up in - `touches_unrefreshed`, not as healthy refreshes, and not as one - traceback per touch.""" - - def touch(h): - raise RuntimeError("backend hiccup") - - worker = _k3_worker(SimpleNamespace(codec=SimpleNamespace(touch=touch))) - ids = [StateTouchId(1, 0), StateTouchId(2, 1), StateTouchId(1, 2)] - with caplog.at_level(logging.WARNING, logger="atom"): - worker._touch_states(SimpleNamespace(state_touches=ids[:2])) - worker._touch_states(SimpleNamespace(state_touches=ids[2:])) - assert worker._touch_verdicts == {i: (True, False) for i in ids} - assert worker._touch_failures == 3 - assert [r.exc_info is not None for r in caplog.records] == [True] - - def _scheduler(self, monkeypatch): - from atom.kv_transfer.offload.dense.connector import DenseOffloadScheduler - from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( - KimiK3OffloadScheduler, - ) - from atom.kv_transfer.offload.metadata import LMCacheOffloadMetadata - - monkeypatch.setattr( - DenseOffloadScheduler, - "build_connector_meta", - lambda self: LMCacheOffloadMetadata(), - ) - s = KimiK3OffloadScheduler.__new__(KimiK3OffloadScheduler) - s._pending_state_loads = [] - s._pending_state_stores = [] - s._pending_state_touches = [] - s._touch_round = 0 - s._state_touch_verdicts = {} - s._state_touches_unrefreshed = 0 - s._refresh_save_stall = lambda: None - return s - - def test_each_touch_is_a_distinct_identity_drained_once(self, monkeypatch): - s = self._scheduler(monkeypatch) - assert s.enqueue_state_touches([]) is False - assert s.enqueue_state_touches([5, 5]) - meta = s.build_connector_meta() - assert meta.state_touches == [StateTouchId(5, 0), StateTouchId(5, 1)] - assert s.build_connector_meta().state_touches == [] - - def test_one_rank_missing_the_image_makes_it_gone(self, monkeypatch): - from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( - STATE_TOUCH_CHANNEL, - ) - - s = self._scheduler(monkeypatch) - for op, ok in ((StateTouchId(5, 0), True), (StateTouchId(5, 1), False)): - handled = s.connector_completion( - ConnectorCompletion(STATE_TOUCH_CHANNEL, op, ok) - ) - assert handled is None, "handled, but not a terminal save" - s.connector_completion( - ConnectorCompletion(STATE_TOUCH_CHANNEL, StateTouchId(6, 2), True) - ) - assert s.take_state_touch_verdicts() == {5: False, 6: True} - assert s.take_state_touch_verdicts() == {} - - def test_a_presence_only_shard_counts_the_touch_unrefreshed(self, monkeypatch): - from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( - STATE_TOUCH_REFRESH_CHANNEL, - ) - - s = self._scheduler(monkeypatch) - for op, ok in ((StateTouchId(5, 0), True), (StateTouchId(6, 1), False)): - handled = s.connector_completion( - ConnectorCompletion(STATE_TOUCH_REFRESH_CHANNEL, op, ok) - ) - assert handled is None, "handled, but not a terminal save" - assert s.take_state_touches_unrefreshed() == 1 - assert s.take_state_touches_unrefreshed() == 0 - assert s.take_state_touch_verdicts() == {}, "not a residency verdict" - - def test_both_touch_reports_reach_quorum_across_tp(self, monkeypatch): - """Each rank files a residency and a refresh report per touch; one - rank that only found the image present counts the touch unrefreshed - without making it gone.""" - from atom.kv_transfer.disaggregation.aggregator import KVOutputAggregator - from atom.kv_transfer.disaggregation.types import KVConnectorOutput - from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( - STATE_TOUCH_CHANNEL, - STATE_TOUCH_REFRESH_CHANNEL, - ) - - touch = StateTouchId(5, 0) - outputs = [] - for result in (TouchResult.REFRESHED, TouchResult.PRESENT): - worker = _k3_worker( - SimpleNamespace(codec=SimpleNamespace(touch=lambda h, r=result: r)) - ) - worker._touch_states(SimpleNamespace(state_touches=[touch])) - resident, refreshed = worker._touch_verdicts[touch] - outputs.append( - KVConnectorOutput( - connector_completions={ - ConnectorCompletion(STATE_TOUCH_CHANNEL, touch, resident), - ConnectorCompletion( - STATE_TOUCH_REFRESH_CHANNEL, touch, refreshed - ), - } - ) - ) - merged = KVOutputAggregator(world_size=2).aggregate(outputs) - s = self._scheduler(monkeypatch) - for completion in merged.connector_completions: - s.connector_completion(completion) - assert s.take_state_touch_verdicts() == {5: True} - assert s.take_state_touches_unrefreshed() == 1 - - def test_the_shells_forward_the_unrefreshed_count(self, monkeypatch): - from atom.kv_transfer.disaggregation.multi.multi_connector import ( - MultiConnectorScheduler, - ) - from atom.kv_transfer.offload.connector import ( - LMCacheOffloadConnectorScheduler, - ) - - shell = LMCacheOffloadConnectorScheduler.__new__( - LMCacheOffloadConnectorScheduler - ) - shell._impl = SimpleNamespace(take_state_touches_unrefreshed=lambda: 2) - assert shell.take_state_touches_unrefreshed() == 2 - shell._impl = SimpleNamespace() - assert shell.take_state_touches_unrefreshed() == 0 - - multi = MultiConnectorScheduler.__new__(MultiConnectorScheduler) - monkeypatch.setattr( - MultiConnectorScheduler, - "_state_tier_sub", - lambda self: SimpleNamespace(take_state_touches_unrefreshed=lambda: 4), - ) - assert multi.take_state_touches_unrefreshed() == 4 - monkeypatch.setattr( - MultiConnectorScheduler, "_state_tier_sub", lambda self: None - ) - assert multi.take_state_touches_unrefreshed() == 0 - - def test_a_negative_round_is_refused(self): - with pytest.raises(ValueError): - StateTouchId(1, -1) - - -# ── the engine: touches published, verdicts applied, admission held ──────── - - -class TestSchedulerPlumbing: - def _sched(self, connector, idx): - from atom.model_engine.scheduler import Scheduler - - s = Scheduler.__new__(Scheduler) - s.kv_connector = connector - s.block_manager = SimpleNamespace(state_offload=idx) - return s - - def test_voted_hashes_reach_the_connector(self, clock): - idx = _index(7) - idx.want_touch(7) - sent = [] - s = self._sched( - SimpleNamespace(enqueue_state_touches=lambda hs: sent.extend(hs) or True), - idx, - ) - s._publish_state_touches() - assert sent == [7] and idx.refresh_hits - - def test_a_connector_without_touches_turns_the_refresh_off(self, clock): - idx = _index(7) - idx.want_touch(7) - s = self._sched(SimpleNamespace(), idx) - s._publish_state_touches() - assert not idx.refresh_hits - clock[0] += STALE - assert not idx.needs_verify(7), "nothing would ever answer the hold" - - def test_verdicts_are_applied_to_the_index(self, clock): - idx = _index(7, 8) - clock[0] += 1 - idx.want_touch(7) - idx.want_touch(8) - assert sorted(idx.take_touches()) == [7, 8] - s = self._sched( - SimpleNamespace(take_state_touch_verdicts=lambda: {7: False, 8: True}), - idx, - ) - s._apply_state_touch_verdicts(idx) - assert not idx.could_serve(7) and idx.could_serve(8) - - def test_presence_only_touches_reach_the_stats_line(self, clock): - idx = _index(7) - s = self._sched( - SimpleNamespace( - take_state_touch_verdicts=lambda: {7: True}, - take_state_touches_unrefreshed=lambda: 3, - ), - idx, - ) - assert idx.stats()["touches_unrefreshed"] == 0 - s._apply_state_touch_verdicts(idx) - s._apply_state_touch_verdicts(idx) - assert idx.stats()["touches_unrefreshed"] == 6 - - -def _paged_bm(): - from test_state_checkpoint import ( - PAGED_COPY_RUNTIME, - ckpt_config, - make_block_manager, - ) - - bm = make_block_manager( - ckpt_config(num_kvcache_blocks=160, pool_entries={"state": 8}), - state_runtime=PAGED_COPY_RUNTIME, - ) - bm.state_offload = StateOffloadIndex(refresh_hits=True) - return bm - - -def _seq(): - from test_state_checkpoint import stateful_seq - - return stateful_seq(list(range(40))) - - -class TestAdmissionHold: - """`record_allocation(defer_unverified=True)` holds an admission that - would load an unconfirmed image, and commits nothing while it does.""" - - def _joint(self, bm, monkeypatch, *, h=77, from_hbm=False): - from atom.model_engine.block_manager import JointBoundaryDecision - - decision = JointBoundaryDecision( - boundary_tokens=16, - boundary_hash=h, - kv_tokens=16, - claim_tokens=16, - state_from_hbm=from_hbm, - ) - commits = [] - monkeypatch.setattr(bm, "_joint_kv_boundary", lambda *a: decision) - monkeypatch.setattr(bm, "_fit_joint", lambda d, probe: d) - monkeypatch.setattr( - bm, "_commit_joint_boundary", lambda seq, d: commits.append(d) - ) - return commits - - def test_an_unconfirmed_joint_image_is_held(self, clock, monkeypatch): - bm = _paged_bm() - commits = self._joint(bm, monkeypatch) - bm.state_offload.note_stored(77) - clock[0] += STALE - seq = _seq() - - assert not bm.record_allocation( - seq, AdmissionProbe(2, [1, 2]), defer_unverified=True - ) - assert commits == [] and bm.state_offload.loads_deferred_for_verify == 1 - assert bm.state_offload.take_touches() == [77] - - bm.state_offload.apply_touch_verdict(77, True) - assert bm.record_allocation( - seq, AdmissionProbe(2, [1, 2]), defer_unverified=True - ) - assert len(commits) == 1 - - def test_a_gone_image_is_not_parked_on(self, clock, monkeypatch): - """The verdict drops the hash; the retry's boundary search no longer - sees the tier rung (`could_serve` is False), so it plans a lower one.""" - bm = _paged_bm() - self._joint(bm, monkeypatch) - bm.state_offload.note_stored(77) - clock[0] += STALE - assert not bm.record_allocation( - _seq(), AdmissionProbe(2, [1, 2]), defer_unverified=True - ) - bm.state_offload.take_touches() - bm.state_offload.apply_touch_verdict(77, False) - assert not bm._tier_can_serve(77) - - def test_a_recently_confirmed_image_goes_straight_through(self, clock, monkeypatch): - bm = _paged_bm() - commits = self._joint(bm, monkeypatch) - bm.state_offload.note_stored(77) - assert bm.record_allocation( - _seq(), AdmissionProbe(2, [1, 2]), defer_unverified=True - ) - assert len(commits) == 1 - - def test_a_state_resident_in_hbm_is_never_held(self, clock, monkeypatch): - bm = _paged_bm() - self._joint(bm, monkeypatch, from_hbm=True) - bm.state_offload.note_stored(77) - clock[0] += STALE - assert bm.record_allocation( - _seq(), AdmissionProbe(2, [1, 2]), defer_unverified=True - ) - - def test_the_fit_probe_path_never_holds(self, clock, monkeypatch): - bm = _paged_bm() - self._joint(bm, monkeypatch) - bm.state_offload.note_stored(77) - clock[0] += STALE - assert bm.record_allocation(_seq(), AdmissionProbe(2, [1, 2])) - - def test_an_hbm_hit_holds_only_for_a_tier_sourced_state(self, clock, monkeypatch): - from atom.model_engine.block_manager import JointBoundaryDecision - - bm = _paged_bm() - monkeypatch.setattr( - bm, - "_joint_kv_boundary", - lambda *a: JointBoundaryDecision(skip_reason="no_kv_reach"), - ) - monkeypatch.setattr(bm, "_commit_joint_boundary", lambda seq, d: None) - bm.state_offload.note_stored(2) - clock[0] += STALE - resident = {2} - monkeypatch.setattr( - bm.paged_state_checkpoints, "contains", lambda h: h in resident - ) - assert bm.record_allocation( - _seq(), AdmissionProbe(2, [1, 2]), defer_unverified=True - ) - resident.clear() - assert not bm.record_allocation( - _seq(), AdmissionProbe(2, [1, 2]), defer_unverified=True - ) - - def test_the_hold_judges_the_joint_fit_kept(self, clock, monkeypatch): - """A joint `_fit_joint` drops is not loaded, so its stale image must not - hold an admission that resumes at an HBM-resident hit.""" - bm = _paged_bm() - commits = self._joint(bm, monkeypatch) - monkeypatch.setattr( - bm, "_fit_joint", lambda d, probe: bm._no_joint("partial_chunk_no_room") - ) - monkeypatch.setattr(bm.paged_state_checkpoints, "contains", lambda h: h == 2) - bm.state_offload.note_stored(77) - clock[0] += STALE - assert bm.record_allocation( - _seq(), AdmissionProbe(2, [1, 2], 1), defer_unverified=True - ) - assert [d.skip_reason for d in commits] == ["partial_chunk_no_room"] - assert bm.state_offload.loads_deferred_for_verify == 0 - - -class TestVotersRegisterTouches: - """`could_serve` only answers; the voters register the hits they count.""" - - def _bm(self): - bm = _paged_bm() - bm.paged_state_checkpoints.attach_offload(bm.state_offload) - bm.state_offload.note_stored(77) - return bm - - def test_the_resumable_scan_registers_a_tier_hit(self, clock): - bm = self._bm() - assert bm.paged_state_checkpoints._reachable(77) - assert bm.state_offload.take_touches() == [77] - - def test_an_hbm_hit_is_not_a_tier_vote(self, clock, monkeypatch): - bm = self._bm() - monkeypatch.setattr( - bm.paged_state_checkpoints.store, "contains", lambda h: True - ) - assert bm.paged_state_checkpoints._reachable(77) - assert bm.state_offload.take_touches() == [] - - def test_asking_whether_the_tier_can_serve_registers_nothing(self, clock): - bm = self._bm() - assert bm._tier_can_serve(77) - assert bm.state_offload.take_touches() == [] - - def test_an_offered_load_registers_its_hash(self, clock): - bm = self._bm() - seq = _seq() - assert bm._request_state_load(seq, 77) - assert bm.state_offload.take_touches() == [77] - - def test_an_offload_without_touches_still_counts_a_tier_hit(self, clock): - """`attach_offload` takes any index answering `could_serve`.""" - bm = _paged_bm() - bm.paged_state_checkpoints.attach_offload( - SimpleNamespace(hashes={77}, could_serve=lambda h: h == 77) - ) - assert bm.paged_state_checkpoints._reachable(77) - assert not bm.paged_state_checkpoints._reachable(78) - - -class TestResumeHash: - """`allocate` and the admission hold pick the state to resume from with - one helper, `_state_resume_hash`.""" - - def test_the_joint_boundary_wins_then_a_replay_then_the_hit(self, monkeypatch): - bm = _paged_bm() - seq = _seq() - assert bm._state_resume_hash(seq, 77, 2) == 77 - assert bm._state_resume_hash(seq, -1, 2) == 2 - assert bm._state_resume_hash(seq, -1, -1) == -1 - monkeypatch.setattr(bm, "state_replay_tokens", 8) - seq.replays_hit = True - assert bm._state_resume_hash(seq, -1, 2) == -1, "a replay restores none" - assert bm._state_resume_hash(seq, 77, 2) == 77 - - def test_a_replayed_hit_is_not_held(self, clock, monkeypatch): - from atom.model_engine.block_manager import JointBoundaryDecision - - bm = _paged_bm() - monkeypatch.setattr( - bm, - "_joint_kv_boundary", - lambda *a: JointBoundaryDecision(skip_reason="no_kv_reach"), - ) - monkeypatch.setattr(bm, "_commit_joint_boundary", lambda seq, d: None) - monkeypatch.setattr(bm.paged_state_checkpoints, "contains", lambda h: False) - monkeypatch.setattr(bm, "state_replay_tokens", 8) - bm.state_offload.note_stored(2) - clock[0] += STALE - seq = _seq() - seq.replays_hit = True - assert bm.record_allocation( - seq, AdmissionProbe(2, [1, 2]), defer_unverified=True - ) - - def test_the_hold_judges_the_hash_allocate_then_loads(self, clock, monkeypatch): - from test_state_checkpoint import boundary_hash, stateful_seq - - bm = _paged_bm() - offload = bm.state_offload - bm.paged_state_checkpoints.attach_offload(offload) - first = _seq() - bm.allocate(first, bm.can_allocate(first)) - bm.hash_blocks(first, bm.checkpoint_limit(first) - first.num_cached_tokens) - h = boundary_hash(bm, first) - bm.take_state_maintenance_ops() - bm.complete_previous_state_batch() - # The image left HBM; only the tier has it, and not confirmed lately. - bm.paged_state_checkpoints.unindex(h) - offload.note_stored(h) - clock[0] += STALE - - asked, attached = [], [] - verify, attach = offload.needs_verify, bm._attach_state_slots - monkeypatch.setattr( - offload, "needs_verify", lambda x: asked.append(x) or verify(x) - ) - monkeypatch.setattr( - bm, "_attach_state_slots", lambda s, x: attached.append(x) or attach(s, x) - ) - second = stateful_seq(list(range(48))) - probe = AdmissionProbe() - probe.num_cached_blocks = bm.can_allocate(second, record=False, probe=probe) - assert probe.num_cached_blocks > 0 - assert not bm.record_allocation(second, probe, defer_unverified=True) - offload.apply_touch_verdict(h, True) - assert bm.record_allocation(second, probe, defer_unverified=True) - assert bm.allocate(second, probe.num_cached_blocks) - assert set(asked) == {h} and attached == [h] - assert second.offload_joint.load_hash == h - - -def test_the_scheduler_retries_a_held_admission(monkeypatch): - """A held seq stays waiting and is admitted on a later pass.""" - from test_state_checkpoint import ckpt_config, make_scheduler, stateful_seq - - sched = make_scheduler(ckpt_config()) - bm = sched.block_manager - real = bm.record_allocation - holds = [True] - - def record(seq, probe, *, defer_unverified=False): - assert defer_unverified, "admission asks to hold unverified loads" - if holds.pop() if holds else False: - return False - return real(seq, probe, defer_unverified=defer_unverified) - - monkeypatch.setattr(bm, "record_allocation", record) - seq = stateful_seq(list(range(12))) - sched.add(seq) - _, scheduled = sched.schedule() - assert scheduled == {} - assert list(sched.waiting) == [seq] and len(seq.block_table) == 0 - _, scheduled = sched.schedule() - assert seq.id in scheduled and not sched.waiting - - -# ── the startup budget check ─────────────────────────────────────────────── - - -GIB = 1 << 30 - - -class TestBudgetWarning: - def test_a_cap_below_half_the_pool_is_fine(self): - assert state_tier_budget_warning(59 * GIB, 208 * GIB, 112 << 20) is None - - def test_a_cap_the_size_of_the_pool_bounds_nothing(self): - """#2562's accuracy run: a 32 GiB pool with the 32 GiB default cap.""" - msg = state_tier_budget_warning(32 * GIB, 32 * GIB, 112 << 20) - assert "bounds nothing" in msg - assert "LMCACHE_MAX_LOCAL_CPU_SIZE=32.0" in msg - - def test_a_cap_over_half_the_pool_is_called_out(self): - msg = state_tier_budget_warning(40 * GIB, 64 * GIB, 112 << 20) - assert msg is not None and "more than half" in msg - - def test_unknown_sizes_are_not_judged(self): - assert state_tier_budget_warning(32 * GIB, 0, 1) is None - assert state_tier_budget_warning(0, 32 * GIB, 1) is None From b1582b46b927f6fb394e9710c0a1ac32c89a709e Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Mon, 2 Nov 2026 03:58:52 +0000 Subject: [PATCH 12/18] state offload: trim the hit refresh to what it needs No behaviour change. - `StateOffloadIndex` reads `OFFLOAD_LOOKUP_DEFER_S` itself (`verify_defer_s` and the test-only `defer_s` argument are gone); `stop_refreshing` is folded into its one caller, `publish_touches`. - The verify window is one constant, 0.25 s per GiB of pool, taken in GiB. - `state_tier_budget_warning` drops the argument only its message used. - `StateTouchId` drops a round check no input can trip, and `enqueue_state_touches` drops an empty-list guard that answered "cannot carry touches" for a call that never happens. - Comments, docstrings and the offload README say the constraint, not the derivation. - Tests: copy-pasted cases parametrized, duplicates merged, function-local imports moved to the top. --- atom/kv_transfer/disaggregation/aggregator.py | 5 +- .../disaggregation/multi/multi_connector.py | 14 +- atom/kv_transfer/disaggregation/types.py | 15 +- atom/kv_transfer/offload/README.md | 35 +-- atom/kv_transfer/offload/connector.py | 9 +- .../offload/hybrid/kimi_k3/connector.py | 44 +-- .../offload/hybrid/kimi_k3/state_object.py | 58 ++-- atom/kv_transfer/offload/metadata.py | 6 +- atom/model_engine/block_manager.py | 25 +- atom/model_engine/page_unit_checkpoint.py | 4 +- atom/model_engine/scheduler.py | 9 +- atom/model_engine/state_offload.py | 255 +++++------------- docs/environment_variables.md | 2 +- tests/test_lmcache_offload_connector.py | 127 +++------ tests/test_state_checkpoint.py | 63 ++--- tests/test_state_offload_index.py | 241 ++++++----------- 16 files changed, 290 insertions(+), 622 deletions(-) diff --git a/atom/kv_transfer/disaggregation/aggregator.py b/atom/kv_transfer/disaggregation/aggregator.py index 7e5da17d44..b614ba030e 100644 --- a/atom/kv_transfer/disaggregation/aggregator.py +++ b/atom/kv_transfer/disaggregation/aggregator.py @@ -157,9 +157,8 @@ 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. + # Touch ids are never reused (fresh round), and tombstoning them would + # push the store / source tombstones out of the ring. self._connector_completions = _TPCompletionGroup[ConnectorCompletionKey]( world_size, terminal_tombstone_limit, diff --git a/atom/kv_transfer/disaggregation/multi/multi_connector.py b/atom/kv_transfer/disaggregation/multi/multi_connector.py index a7cefe9ee5..0d4a6f6ae7 100644 --- a/atom/kv_transfer/disaggregation/multi/multi_connector.py +++ b/atom/kv_transfer/disaggregation/multi/multi_connector.py @@ -656,23 +656,15 @@ def enqueue_state_stores(self, stores) -> bool: 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`. - """ + """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. - - Select by `has_state_tier`, not method presence -- see - `_state_tier_sub`. - """ + """`{hash: resident}` from the tier sub; empty if none.""" c = self._state_tier_sub() if c is None: return {} diff --git a/atom/kv_transfer/disaggregation/types.py b/atom/kv_transfer/disaggregation/types.py index b46edb1036..177041a521 100644 --- a/atom/kv_transfer/disaggregation/types.py +++ b/atom/kv_transfer/disaggregation/types.py @@ -126,23 +126,14 @@ def __post_init__(self) -> None: @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. + """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. `round` keeps repeated touches of one hash distinct identities. """ prefix_hash: int round: int - def __post_init__(self) -> None: - if self.round < 0: - raise ValueError("state touch round must be nonnegative") - ConnectorCompletionId = ( ReqId diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index d50b9c3277..fd534bf0b8 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -122,27 +122,18 @@ 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 touch that raises keeps the -image indexed and is counted there too (one traceback per worker). 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`. +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 +`state_offload_touches_unrefreshed`. `OFFLOAD_STATE_CPU_SIZE` caps the images' +share of the pool and reserves nothing; the worker warns when it is not below +half of the pool. ## LMCache Multiprocess (`lmcache_mp`) @@ -1467,7 +1458,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` | +| `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) diff --git a/atom/kv_transfer/offload/connector.py b/atom/kv_transfer/offload/connector.py index 352c0fa697..dacf39499a 100644 --- a/atom/kv_transfer/offload/connector.py +++ b/atom/kv_transfer/offload/connector.py @@ -336,20 +336,17 @@ def enqueue_state_stores(self, stores) -> bool: 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).""" + """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; empty when the - impl carries none.""" + """`{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_touches_unrefreshed(self) -> int: - """Touches that only confirmed presence; 0 when the impl carries none.""" + """Touches that only confirmed presence.""" callback = getattr(self._impl, "take_state_touches_unrefreshed", None) return int(callback()) if callback is not None else 0 diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py index 647c41ea3d..049fb737d1 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py @@ -68,14 +68,12 @@ #: rather than leaving a partially-reported key pending forever. STATE_LOAD_DISPOSITION_CHANNEL = "k3_state_load_disposition" -#: One verdict per `StateTouchId`: the worker refreshed the image in LMCache's -#: LRU (`succeeded=True`) or found it gone (`succeeded=False`). Failure-dominant -#: across TP, like a joint load: one missing shard means the load would miss. +#: 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" -#: Alongside each `STATE_TOUCH_CHANNEL` verdict, same id: `succeeded=False` when -#: the image was present but outside LocalCPU's hot cache, so the touch did not -#: refresh its eviction order. Failure-dominant: any such shard counts. +#: 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 @@ -267,7 +265,7 @@ def _build_state_tier(self, transfer_tensors) -> None: codec.bind_storage_manager(self._engine.storage_manager) pool_gib = float(getattr(cfg, "max_local_cpu_size", 0) or 0) budget_warning = state_tier_budget_warning( - state_bytes, int(pool_gib * (1 << 30)), entry_bytes + state_bytes, int(pool_gib * (1 << 30)) ) if budget_warning: logger.warning("kimi_k3 offload: %s", budget_warning) @@ -314,14 +312,10 @@ def start_load_kv(self, metadata) -> None: super().start_load_kv(metadata) def _touch_states(self, metadata) -> None: - """Refresh the state images the engine voted a hit for, and record - whether each is still there. - - Inline on this thread, not on the tier's executors: a touch is one - dict move under LocalCPU's lock, and it must not queue behind an - image-sized store. With no tier every image is gone by definition, so - the engine's index drops what it still advertises. - """ + """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 @@ -333,9 +327,8 @@ def _touch_states(self, metadata) -> None: try: result = tier.codec.touch(int(touch_id.prefix_hash)) except Exception: - # Says nothing about the bytes, so the index keeps the image; - # nothing was refreshed either. One traceback: a drifted - # LMCache raises on every touch. + # Says nothing about the bytes: kept, not refreshed. One + # traceback, since a drifted LMCache raises on every touch. self._touch_failures += 1 if self._touch_failures == 1: logger.warning( @@ -682,13 +675,9 @@ 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] = [] - # Hit refresh: touches waiting for the next metadata, the round - # counter that keeps each one a distinct aggregator identity, and the - # verdicts drained by the engine (`take_state_touch_verdicts`). self._pending_state_touches: list[StateTouchId] = [] self._touch_round = 0 self._state_touch_verdicts: dict[int, bool] = {} - # Touches that only confirmed presence (`take_state_touches_unrefreshed`). 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 @@ -728,13 +717,7 @@ def enqueue_state_stores(self, stores) -> bool: return True def enqueue_state_touches(self, hashes) -> bool: - """Queue LRU touches for state images the engine just voted a hit for. - - Each gets its own round, so a hash touched again is a new aggregator - identity rather than a tombstoned duplicate. - """ - if not hashes: - return False + """Queue touches for the next metadata, each under its own round.""" for h in hashes: self._pending_state_touches.append(StateTouchId(int(h), self._touch_round)) self._touch_round += 1 @@ -747,8 +730,7 @@ def take_state_touch_verdicts(self) -> dict[int, bool]: return verdicts def take_state_touches_unrefreshed(self) -> int: - """Touches since the last call that found an image present but could - not refresh it in LocalCPU's LRU.""" + """Touches since the last call that only confirmed presence.""" n = self._state_touches_unrefreshed self._state_touches_unrefreshed = 0 return n diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py index e907f6850e..90ce8e303a 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py @@ -26,24 +26,16 @@ class TouchResult(enum.IntEnum): """What a touch found. Truthy when the image can still be read.""" GONE = 0 - #: Outside LocalCPU's hot cache but held by another backend: present, its - #: eviction order not refreshed. `get` still reads it, through LMCache's - #: `StorageManager.get`, which writes a hit from another backend back into - #: LocalCPU (`LocalCPUBackend.submit_put_task`); nothing in ATOM does. + #: Held only by another backend: present, not refreshed. `get` still reads + #: it; LMCache's `StorageManager.get` writes it back into LocalCPU. PRESENT = 1 REFRESHED = 2 def _local_cpu_lru(storage_manager): - """LMCache's LocalCPU backend, checked once for what a touch needs, or - None when a touch cannot refresh anything. - - `cpu_lock`, `hot_cache` and `cache_policy.update_on_hit` are LMCache - internals (written against 0.5.6), not API. Without one of them, touches - only confirm presence, like a storage manager with no LocalCPU backend: - loads still work, only the refresh is lost. Not silent: it warns here and - every touch is counted in `state_offload_touches_unrefreshed`. - """ + """LMCache's LocalCPU backend if it has the internals a touch uses + (`cpu_lock`, `hot_cache`, `cache_policy.update_on_hit`; LMCache 0.5.6), + else None, warned: touches then only confirm presence.""" backend = getattr(storage_manager, "local_cpu_backend", None) if backend is None: missing = ["local_cpu_backend"] @@ -114,9 +106,7 @@ def __init__( ) if self._max_entries is not None and self._max_entries < 1: raise ValueError("state CPU budget must fit at least one checkpoint") - # The cap's LRU, in the same order as LMCache's: refreshed by `get`, - # `touch` and a re-store, so the cap retires images nobody resumes - # from rather than the oldest stores. + # The cap's LRU, kept in LMCache's order (`get`, `touch`, a re-store). self._cache_keys: OrderedDict[Any, None] = OrderedDict() self._cache_lock = threading.Lock() self._store_lock = threading.Lock() @@ -169,19 +159,15 @@ def key(self, h: int): ) def put(self, h: int, unit_ids, on_source_released=None) -> bool: - # Cap state independently of KV while retaining the shared allocator; - # within the cap, images and KV chunks compete in LMCache's one LRU. + # Bound state independently of KV while retaining the shared allocator. # Only stores serialize; a load holds its MemoryObj reference while # unpacking, so removing its cache entry cannot free live bytes. if self._storage is None: return False key = self.key(h) if self._touch_key(key) is TouchResult.REFRESHED: - # Same prefix hash and layout, same bytes, still in LocalCPU's hot - # cache: storing again would only cost an image-sized D2H that - # LocalCPU then refuses as a duplicate key. The touch counts as the - # use it is. Only there: an image another backend holds (disk) is - # stored again, which puts it back into the CPU tier. + # Same key, same bytes, already hot: a re-store is a D2H LocalCPU + # refuses as a duplicate. One held only by disk is stored again. if on_source_released is not None: on_source_released() return True @@ -202,27 +188,18 @@ def put(self, h: int, unit_ids, on_source_released=None) -> bool: return ok def touch(self, h: int) -> TouchResult: - """Refresh `h`'s image in LMCache's LRU; falsy (`GONE`) when it is gone. - - The state-image twin of what a KV lookup does to its chunks - (`LocalCPUBackend.batched_async_contains(pin=True)` then - `touch_cache()` -> `cache_policy.update_on_hit`). State hits are - answered engine-side by `StateOffloadIndex`, so without this LMCache - never learns an image was used, and a prefix hit on every turn keeps - its KV at the young end and its image at the cold end, where pressure - evicts it first. - - `update_on_hit` under the backend's `cpu_lock`, not `contains(pin=True)` - + `touch_cache()`: a pin needs a matching unpin, and `keys_in_request` - is the list the KV lookup is filling at the same time. + """Refresh `h`'s image in LMCache's LRU, as a KV lookup does for its + chunks; falsy (`GONE`) when it is gone. A state hit is answered by + `StateOffloadIndex`, so otherwise LMCache never sees the image used. + + `update_on_hit` under `cpu_lock`, not `contains(pin=True)` + + `touch_cache()`: a pin needs an unpin, and `keys_in_request` belongs + to the concurrent KV lookup. """ if self._storage is None: return TouchResult.GONE result = self._touch_key(self.key(h)) if result is TouchResult.PRESENT: - # Still reported resident: `get` can read it now, and the - # confirmation that opens only lets later loads skip the hold, - # where a miss degrades to a recompute. self._unrefreshed_touches += 1 if self._unrefreshed_touches == 1: logger.warning( @@ -251,8 +228,7 @@ def _refresh(self, key) -> TouchResult: if key in backend.hot_cache: backend.cache_policy.update_on_hit(key, backend.hot_cache) return TouchResult.REFRESHED - # Not in the hot cache (or there is none): another backend may still - # hold it, and `get` reads every backend. Presence only. + # Another backend may still hold it, and `get` reads every backend. if self._storage.contains(key): return TouchResult.PRESENT return TouchResult.GONE diff --git a/atom/kv_transfer/offload/metadata.py b/atom/kv_transfer/offload/metadata.py index d54a40498c..5cfab3aaf2 100644 --- a/atom/kv_transfer/offload/metadata.py +++ b/atom/kv_transfer/offload/metadata.py @@ -216,10 +216,8 @@ def __init__(self) -> None: # generation) is what the source-release and index reports carry back to # settle the pin and index the hash. self.state_stores: list[tuple] = [] - # StateTouchId per state image the engine voted a hit for: each worker - # refreshes the image in LMCache's LRU and reports whether it is still - # there. Work in its own right -- an admission can be waiting on the - # verdict. + # StateTouchId per voted state image, refreshed and reported by each + # worker. self.state_touches: list = [] def add_request(self, meta: LMCacheReqMeta) -> None: diff --git a/atom/model_engine/block_manager.py b/atom/model_engine/block_manager.py index 8ed11fc226..0d3ecd84ad 100644 --- a/atom/model_engine/block_manager.py +++ b/atom/model_engine/block_manager.py @@ -516,10 +516,7 @@ def __init__( if checkpoint_spec is not None else None ), - # Off again (`stop_refreshing`) if the connector turns out not - # to carry state touches. refresh_hits=True, - # How long a confirmed image is trusted. verified_for_s=state_tier_verified_for_s(config), ) # Attached rather than passed at construction: the coordinator is @@ -1582,12 +1579,8 @@ def _resume_hash_off_hbm( block_hashes: list[int], decision: JointBoundaryDecision, ) -> int: - """The hash `allocate` would restore this seq's state from when HBM - does not hold its state, or -1 (no state, HBM holds it, or a replay). - - The hash `allocate` will hand `_attach_state_slots` - (`_state_resume_hash`), so a hold judges what will be loaded. - """ + """The `_state_resume_hash` `allocate` would restore from, when HBM + does not hold it; -1 for no state, an HBM-held one, or a replay.""" if not seq.has_per_req_cache: return -1 joint = decision.skip_reason is None @@ -1602,10 +1595,12 @@ def _resume_hash_off_hbm( seq, decision.boundary_hash if joint else -1, hit_hash ) if h != -1 and not joint: - if self.paged_state_checkpoints is not None: - if self.paged_state_checkpoints.contains(h): - return -1 - elif self.state.lookup(h) >= 0: + in_hbm = ( + self.paged_state_checkpoints.contains(h) + if self.paged_state_checkpoints is not None + else self.state.lookup(h) >= 0 + ) + if in_hbm: return -1 return h @@ -1614,9 +1609,7 @@ def _state_resume_hash(self, seq: Sequence, joint_hash: int, hit_hash: int) -> i The joint boundary when there is one (above the HBM hit by construction), else the HBM hit's last block. A replay rebuilds the - state instead of restoring it, so it resumes from none. `allocate` - hands this to `_attach_state_slots`, and `_resume_hash_off_hbm` reports - it, so the admission hold judges the hash that will be loaded. + state instead of restoring it, so it resumes from none. """ if joint_hash != -1: return joint_hash diff --git a/atom/model_engine/page_unit_checkpoint.py b/atom/model_engine/page_unit_checkpoint.py index 2adb966c81..357c8a8380 100644 --- a/atom/model_engine/page_unit_checkpoint.py +++ b/atom/model_engine/page_unit_checkpoint.py @@ -1072,9 +1072,7 @@ def _reachable(self, h: int) -> bool: `hashes` from its own stores but cannot load them back, and voting a hit it would then refuse strands the scan at the tier rung -- skipping a still-resident HBM rung -- and forfeits the checkpoint to a recompute. - - A tier hit accepted here is a vote, so it is registered for a touch - (`want_touch`). + A tier hit accepted here is a vote (`want_touch`). """ if self.store.contains(h): return True diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index 8965792dd2..9159154ff7 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -2382,10 +2382,9 @@ def _schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]] | None: if not self.block_manager.record_allocation( seq, hbm_probe, defer_unverified=True ): - # The state offload holds it a pass (the state it would load - # is unconfirmed, and its check is out). Held like an - # in-flight prefix: an image found gone is forgotten when the - # verdict lands, and the retry plans without it. + # Held a pass for the state offload's verdict on the image it + # would load; a gone image is forgotten and the retry plans + # without it. prefix_waiters.append(seq) continue if not self.block_manager.allocate(seq, num_cached_blocks): @@ -4918,7 +4917,7 @@ def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput): bm = getattr(self, "block_manager", None) offload = getattr(bm, "state_offload", None) if offload is not None: - # Before the next admission pass, which reads what they decide. + # Before the next admission pass reads them. offload.collect_touch_verdicts(self.kv_connector) take = getattr(self.kv_connector, "take_state_reports", None) if offload is not None and take is not None: diff --git a/atom/model_engine/state_offload.py b/atom/model_engine/state_offload.py index e27dc4e23c..2e1c405639 100644 --- a/atom/model_engine/state_offload.py +++ b/atom/model_engine/state_offload.py @@ -17,20 +17,17 @@ from dataclasses import dataclass, replace from time import monotonic +from atom.utils import envs + logger = logging.getLogger("atom") -#: The fastest one worker's LMCache pool is assumed to take in new entries. -#: AgentX c64 on one GPU stored about 0.25 GiB/s (970 GB/h); four times that -#: leaves room for a prefill-heavier load or a model with more KV per token. -_POOL_FILL_GIB_PER_S = 1.0 -#: The share of the pool's turnover at that rate an image is trusted for. -_POOL_TRUSTED_FRACTION = 0.25 -#: Bounds on `tier_verified_for_s`. See there. +#: Seconds a confirmed image is trusted per GiB of LMCache pool: a quarter of +#: the pool's turnover at 1 GiB/s, four times the fill rate AgentX c64 measured. +_TIER_VERIFIED_S_PER_GIB = 0.25 _TIER_VERIFIED_MIN_S = 5.0 _TIER_VERIFIED_MAX_S = 60.0 -#: A hash voted on every scheduling pass is touched at most this often. The -#: touch only has to keep the image near the young end, not track every pass. +#: A hash voted on every scheduling pass is touched at most this often. _TIER_TOUCH_INTERVAL_S = 0.5 @@ -41,14 +38,11 @@ class StateOffloadIndex: can drop bytes at any time, so on its own a hash here means "was stored once", never "is still there". - With `refresh_hits`, every hit a voter registers (`want_touch`) is also - *touched*: the hash is queued for `take_touches`, the worker refreshes the image's - position in LMCache's LRU the same way a KV lookup refreshes its chunks, - and reports whether the image was still there (`apply_touch_verdict`). That - keeps a prefix's state image as warm as its KV, and lets a hash LMCache - dropped leave the index before a request parks on it (`needs_verify`). - A false positive that slips through still costs one park/unpark and is - handled by the `failed_loading` path (`fail_load` -> `forget`). + With `refresh_hits`, every voted hit is also touched: the worker refreshes + the image in LMCache's LRU, as a KV lookup does for its chunks, and reports + whether it was still there (`apply_touch_verdict`). A false positive that + slips through still costs one park/unpark and is handled by the + `failed_loading` path (`fail_load` -> `forget`). """ def __init__( @@ -58,8 +52,7 @@ def __init__( can_load: bool = True, max_cpu_entries: int | None = None, refresh_hits: bool = False, - verified_for_s: float | None = None, - defer_s: float | None = None, + verified_for_s: float = _TIER_VERIFIED_MIN_S, ) -> None: # The two legs are separately granted. A `kv_producer` role saves and # never loads, a `kv_consumer` loads and never saves, and offering a @@ -111,44 +104,26 @@ def __init__( # than indexed. Any non-zero value means the reclaim window is firing # on live transfers and is set too low. self.stores_untrusted = 0 - # Hit refresh (see the class docstring). Off for indexes whose - # connector carries no touches (the vLLM plugin), where the queue would - # only grow. + # Off for indexes whose connector carries no touches (the vLLM plugin). self.refresh_hits = bool(refresh_hits) - # Both fixed for the server's lifetime; `needs_verify` runs on the - # admission path. The window comes from the tier's pool - # (`state_tier_verified_for_s`); the floor when not given. - self.verified_for_s = ( - tier_verified_for_s(0) if verified_for_s is None else float(verified_for_s) - ) - if defer_s is None: - defer_s = verify_defer_s() if self.refresh_hits else 0.0 - self.defer_s = float(defer_s) - # Hashes voted for since the last `take_touches`. + self.verified_for_s = float(verified_for_s) + # Read once: `needs_verify` runs on the admission path. Not read + # without the refresh, which never holds. + self.defer_s = float(envs.OFFLOAD_LOOKUP_DEFER_S) if self.refresh_hits else 0.0 self._touch_wanted: set[int] = set() - # hash -> monotonic of the last touch handed out / the last evidence - # the image was resident. + # hash -> monotonic of the last touch sent / the last evidence the + # image was resident. self._touch_sent_at: dict[int, float] = {} self._verified_at: dict[int, float] = {} - # hash -> monotonic of the first admission held for its verdict; the - # `defer_s` bound runs from here, so re-sent touches cannot extend it. - # Cleared by any evidence (verdict, store, load) or by dropping `h`. + # hash -> monotonic of the first held admission; `defer_s` runs from + # here, so re-sent touches cannot extend a hold. self._hold_started_at: dict[int, float] = {} - # Hashes whose hold ran out: admissions over them go ahead unheld - # until evidence arrives. + # Holds that ran out; admitted unheld until evidence arrives. self._hold_expired: set[int] = set() self.touches_sent = 0 - # Touches that found the image gone: LMCache evicted it and the index - # learned of it before any request parked on it. self.touch_misses = 0 - # Touches that found the image present but outside LocalCPU's hot - # cache, so they confirmed it without refreshing its eviction order. self.touches_unrefreshed = 0 - # Admissions held back a pass because the image they would load had - # not been confirmed recently (`needs_verify`). self.loads_deferred_for_verify = 0 - # Holds released because no verdict came back within `defer_s`. Any - # non-zero value means touch verdicts are being lost. self.verify_holds_expired = 0 # ------------------------------- stores -------------------------------- # @@ -194,8 +169,8 @@ def _end_hold(self, h: int) -> None: def could_serve(self, h: int) -> bool: """Whether a load for `h` could be offered at all: this tier can load - AND believes it holds `h`. Pure; a voter that acts on a True answer - registers the hit with `want_touch`. + AND believes it holds `h`. Pure: a voter registers its hit with + `want_touch`. The membership+capability half of `request_load`'s guard (minus its per-request in-flight check), factored out so the admission-path voters @@ -212,22 +187,14 @@ def could_serve(self, h: int) -> bool: # ------------------------------- touches ------------------------------- # def want_touch(self, h: int) -> None: - """A voter counted `h` as a hit: under `refresh_hits`, queue it for - `take_touches`. - - Called on every hit voted for (the resumable scan, and `request_load` - for the load it offers), so a prefix whose state is voted - for on every pass stays as warm in LMCache as its KV chunks, which the - KV lookup refreshes on every pass. - """ + """A voter counted `h` as a hit (the resumable scan, or `request_load`); + under `refresh_hits`, queue it for `take_touches`.""" if self.refresh_hits and h in self.hashes: self._touch_wanted.add(h) def take_touches(self) -> list[int]: - """Hashes to touch in LMCache this step, drained. - - A hash already touched within `_TIER_TOUCH_INTERVAL_S` is skipped: a - request waiting for admission votes for the same hash on every pass. + """Hashes to touch this step, drained; one touched within + `_TIER_TOUCH_INTERVAL_S` is skipped (a waiting request votes every pass). """ if not self._touch_wanted: return [] @@ -243,19 +210,12 @@ def take_touches(self) -> list[int]: return out def apply_touch_verdict(self, h: int, resident: bool) -> None: - """A touch came back: the image was still in LMCache, or it was not. - - Gone means LMCache's LRU dropped it, so the hash leaves the index now, - before a request parks a joint load on it and recomputes the whole - prompt after the miss. - """ + """A touch came back; gone drops `h` before a request parks on it.""" h = int(h) if h not in self.hashes: return if not resident: - # Confirmed after the latest touch was sent (its store landed since, - # or the hash was forgotten and stored again): the touch looked - # before those bytes were there, so its verdict is stale. + # Stale: `h` was stored (again) after this touch looked. if self._verified_at.get(h, -1e18) > self._touch_sent_at.get(h, -1e18): return self.touch_misses += 1 @@ -265,24 +225,10 @@ def apply_touch_verdict(self, h: int, resident: bool) -> None: # The touch made it LMCache's youngest image; keep the cap's order. self._hash_lru.move_to_end(h) - def stop_refreshing(self) -> None: - """The connector cannot carry touches: stop queueing them and stop - holding admissions for verdicts that will never come.""" - self.refresh_hits = False - self._touch_wanted.clear() - self._touch_sent_at.clear() - self._hold_started_at.clear() - self._hold_expired.clear() - def publish_touches(self, connector) -> None: - """Hand this pass's touches to `connector`, which refreshes the images - in LMCache's LRU and reports whether each is still there - (`collect_touch_verdicts`). - - A connector that will not carry them turns the refresh off for good: - otherwise every admission that waits for a verdict would wait out - `defer_s` for nothing. - """ + """Hand this pass's touches to `connector`. One that will not carry + them turns the refresh off for good, so no admission waits out + `defer_s` for a verdict that never comes.""" if not self.refresh_hits: return hashes = self.take_touches() @@ -297,11 +243,14 @@ def publish_touches(self, connector) -> None: "evicts is only noticed when a load misses.", type(connector).__name__, ) - self.stop_refreshing() + self.refresh_hits = False + self._touch_wanted.clear() + self._touch_sent_at.clear() + self._hold_started_at.clear() + self._hold_expired.clear() def collect_touch_verdicts(self, connector) -> None: - """Apply the touch verdicts `connector` reached TP quorum on, and add - its presence-only touches to `touches_unrefreshed`.""" + """Apply the verdicts `connector` reached TP quorum on.""" take = getattr(connector, "take_state_touch_verdicts", None) if take is None: return @@ -312,27 +261,17 @@ def collect_touch_verdicts(self, connector) -> None: self.touches_unrefreshed += int(unrefreshed()) def needs_verify(self, h: int) -> bool: - """Whether an admission about to load `h` should wait a pass for a - touch verdict instead of parking on it now. - - True while `h` has not been confirmed resident within - `verified_for_s` and its touch is still pending. A pass later - the verdict is normally in: resident lets the load go ahead, gone has - already dropped `h`, so the admission picks a lower boundary (an HBM - checkpoint, an older CPU image, or a plain prefill) instead of parking - and recomputing everything after the miss. - - The hold is bounded by `defer_s` (`OFFLOAD_LOOKUP_DEFER_S`, read once - at construction) from the first time `h` was held, not - from its latest touch: a held request re-probes every pass and its - touch is re-sent every `_TIER_TOUCH_INTERVAL_S`, so a bound measured - from the send would never run out if verdicts stopped arriving. Once - it runs out, admissions over `h` proceed as they did before touches - existed until any evidence about `h` arrives. + """Whether an admission about to load `h` waits a pass for its touch + verdict: `h` not confirmed within `verified_for_s`. If the verdict is + gone, `h` has left the index and the retry plans a lower boundary + instead of parking on a miss. + + Bounded by `defer_s` from the first hold, not from the latest touch + (re-sent every `_TIER_TOUCH_INTERVAL_S`, so it would never run out); + after that `h` is admitted unheld until new evidence arrives. """ defer_s = self.defer_s - # `could_serve`, not `h in hashes`: a store-only role (`kv_producer`) - # indexes its own stores but never loads them. + # `could_serve`: a store-only role indexes stores it never loads. if not self.refresh_hits or defer_s <= 0 or not self.could_serve(h): return False now = monotonic() @@ -359,9 +298,8 @@ def needs_verify(self, h: int) -> bool: return True def hold_admission(self, h: int) -> bool: - """Whether an admission that would restore its state from `h` (-1: - from nothing outside HBM) waits a pass for `h`'s touch verdict - (`needs_verify`). Counted in `loads_deferred_for_verify`.""" + """`needs_verify` for an admission restoring its state from `h` (-1: + nothing outside HBM), counted.""" if h == -1 or not self.needs_verify(h): return False self.loads_deferred_for_verify += 1 @@ -373,8 +311,8 @@ def request_load(self, req_id, h: int) -> bool: The guard between believing and delivering: a load is resolved only by a worker report, so offering one for a hash never stored would park the - request against bytes no `get` can produce. An offered load is a hit, - so it is registered for a touch (`want_touch`). + request against bytes no `get` can produce. An offered load is a hit + (`want_touch`). """ if not self.could_serve(h): # A store-only role (can_load False) or a hash never stored. Voting @@ -462,11 +400,11 @@ def stats(self) -> dict[str, int]: # far past `indexed` means the state budget is smaller than the # reusable working set. "hashes_evicted": self.hashes_evicted, - # Hit refresh. `touch_misses` are images LMCache dropped that left - # the index before anything parked on them; `loads_failed` is what - # still got past. `touches_unrefreshed` is the part of the - # resident verdicts that only confirmed presence: 0 when every - # touch refreshed LMCache's LRU. + # Hit refresh. `touch_misses`: images LMCache dropped that left the + # index before anything parked on them (`loads_failed` is what + # still got past). `touches_unrefreshed`: confirmed present but + # not refreshed. `verify_holds_expired`: holds that got no verdict + # within `defer_s`; non-zero means verdicts are being lost. "touches_sent": self.touches_sent, "touch_misses": self.touch_misses, "touches_unrefreshed": self.touches_unrefreshed, @@ -475,46 +413,15 @@ def stats(self) -> dict[str, int]: } -def verify_defer_s() -> float: - """How long an admission may wait for a touch verdict - (`StateOffloadIndex.needs_verify`), read once per index. +def tier_verified_for_s(pool_gib: float) -> float: + """How long an image counts as resident after the last evidence (store, + load, touch), which left it at the young end of LMCache's LRU. - `OFFLOAD_LOOKUP_DEFER_S`: the same question it already answers for a KV - lookup in flight -- how long admission may pass over a request whose tier - answer is on its way rather than block on it or guess. + A share of the pool's turnover, clamped: the floor stays well above a + touch round trip, the ceiling is the largest window measured (AgentX c48 / + c64). Either error costs a held pass or a missed load, not wrong output. """ - from atom.utils import envs - - return float(envs.OFFLOAD_LOOKUP_DEFER_S) - - -def tier_verified_for_s(pool_bytes: int) -> float: - """How long a hash counts as known-resident after the last evidence that - its image was in LMCache (stored, loaded, or confirmed by a touch). - - Each of those leaves the image at the young end of LMCache's LRU, which - evicts it only after a whole pool of newer entries has come in: the safe - window is the pool's turnover, `pool_bytes / fill rate`. The fill rate is - the workload's and not known at startup, so this takes a ceiling, - `_POOL_FILL_GIB_PER_S`, and trusts `_POOL_TRUSTED_FRACTION` of the - turnover at that rate: 0.25 s per GiB, 52 s for the 208.6 GiB AgentX pool, - whose turnover at c64 was about 14 minutes. - - Clamped to [`_TIER_VERIFIED_MIN_S`, `_TIER_VERIFIED_MAX_S`] = [5 s, 60 s]. - The floor (pools under 20 GiB) keeps the window well above a touch round - trip, `_TIER_TOUCH_INTERVAL_S` plus a step, or nearly every load would be - held. The ceiling (pools over 240 GiB) is the largest window measured - (AgentX c48 / c64); past it a fill rate above the assumed one would let - images evicted minutes earlier through. An unknown pool (0) gets the floor. - - Erring either way is a cost, not wrong output: too long admits a load that - may miss and recompute, which is what every load did before touches; too - short holds a load a pass for a verdict it did not need. - """ - if pool_bytes <= 0: - return _TIER_VERIFIED_MIN_S - turnover_s = pool_bytes / float(1 << 30) / _POOL_FILL_GIB_PER_S - window = turnover_s * _POOL_TRUSTED_FRACTION + window = pool_gib * _TIER_VERIFIED_S_PER_GIB return min(max(window, _TIER_VERIFIED_MIN_S), _TIER_VERIFIED_MAX_S) @@ -692,25 +599,15 @@ def state_tier_cpu_bytes() -> int: return int(float(os.environ.get("OFFLOAD_STATE_CPU_SIZE", "32")) * (1 << 30)) -def state_tier_budget_warning( - state_bytes: int, pool_bytes: int, entry_bytes: int -) -> str | None: - """Why the state cap and the LMCache pool size do not fit together, or None. - - The cap only bounds how much of the shared pool images may take, so a cap - at or above the pool bounds nothing, and one above half of it lets images - push most of the KV out under pressure. Neither stops a store -- the pool - evicts by recency either way -- but both are almost always a sizing - mistake (e.g. `LMCACHE_MAX_LOCAL_CPU_SIZE=32` with the 32 GiB default cap), - so the caller says so loudly at startup. - """ +def state_tier_budget_warning(state_bytes: int, pool_bytes: int) -> str | None: + """Why the state cap and the LMCache pool it shares do not fit together + (e.g. `LMCACHE_MAX_LOCAL_CPU_SIZE=32` with the 32 GiB default cap), or + None.""" if pool_bytes <= 0 or state_bytes <= 0: return None gib = float(1 << 30) sizes = ( - f"OFFLOAD_STATE_CPU_SIZE={state_bytes / gib:.1f} GiB " - f"({state_bytes // max(1, entry_bytes)} images of " - f"{entry_bytes / (1 << 20):.1f} MiB), " + f"OFFLOAD_STATE_CPU_SIZE={state_bytes / gib:.1f} GiB, " f"LMCACHE_MAX_LOCAL_CPU_SIZE={pool_bytes / gib:.1f} GiB" ) if state_bytes >= pool_bytes: @@ -781,15 +678,11 @@ def state_tier_chunk_tokens(config) -> int: def state_tier_verified_for_s(config) -> float: - """`tier_verified_for_s` for one worker's LMCache CPU pool - (`LMCACHE_MAX_LOCAL_CPU_SIZE`): how long `StateOffloadIndex` trusts a - confirmed image. The floor when LMCache's config cannot be read. - - Like `state_tier_chunk_tokens`, call it only when a tier is hosted. - """ + """`tier_verified_for_s` for one worker's LMCache CPU pool, the floor + when it cannot be read. Like `state_tier_chunk_tokens`, call it only when + a tier is hosted.""" try: - cfg = _lmcache_engine_config(config) - pool_gib = float(getattr(cfg, "max_local_cpu_size", 0) or 0) + pool_gib = float(_lmcache_engine_config(config).max_local_cpu_size) except Exception: # noqa: BLE001 - `state_tier_chunk_tokens` logged why logger.warning( "state offload: LMCache's CPU pool size is unknown; a confirmed " @@ -797,7 +690,7 @@ def state_tier_verified_for_s(config) -> float: _TIER_VERIFIED_MIN_S, ) pool_gib = 0.0 - return tier_verified_for_s(int(pool_gib * (1 << 30))) + return tier_verified_for_s(pool_gib) #: The connector backends whose worker half can build a `StateOffloadTier`. diff --git a/docs/environment_variables.md b/docs/environment_variables.md index 50fd6881d5..87d299f658 100644 --- a/docs/environment_variables.md +++ b/docs/environment_variables.md @@ -413,7 +413,7 @@ derive a bound. | **OFFLOAD_LOOKUP_MEMO_STEPS** | int | 32 | Scheduler steps a memoised tier-lookup answer is replayed before the tier is asked again. | | **OFFLOAD_LOOKUP_RETRY_STEPS** | int | 32 | Scheduler steps a failed tier lookup suppresses the next attempt. | | **OFFLOAD_ASYNC_LOOKUP_DEPTH** | int | 16 | `lmcache_mp` sends tier lookups for the head of the waiting queue before admission, as non-blocking requests, and admission consumes the answer; a request whose answer has not arrived is passed over (in arrival order) instead of blocking the scheduler. This is how many waiting requests per scheduling pass get a prefetched lookup. | -| **OFFLOAD_LOOKUP_DEFER_S** | float | 2.0 | How long `lmcache_mp` admission may pass over a request whose prefetched lookup is still in flight before it waits for the answer (finite, ≥ 0). Also bounds how long `lmcache_offload` with a state tier (Kimi-K3, Qwen) holds an admission that would load a CPU state image not confirmed recently, while that image's LRU touch reports whether LMCache still holds it; `0` disables the hold. | +| **OFFLOAD_LOOKUP_DEFER_S** | float | 2.0 | How long `lmcache_mp` admission may pass over a request whose prefetched lookup is still in flight before it waits for the answer (finite, ≥ 0). Also bounds how long a `lmcache_offload` state tier holds an admission for the LRU touch verdict on the CPU state image it would load; `0` disables the hold. | | **OFFLOAD_SKIP_LOOKUP_ON_HBM_HIT** | bool | 1 (true) | On by default; set to `0` to always wait for the lookup. Admit a prefill without waiting for its external-tier lookup when the lookup cannot change what is loaded: the HBM prefix leaves fewer than `OFFLOAD_MIN_LOAD_TOKENS` tokens uncached, so every load path would refuse the tier hit. Requires an asynchronous lookup (LMCache MP) and is not applied with joint KV/state offload. The lookup is still sent, because it refreshes the prefix's L1 LRU recency and opens the MP session whose end touches the request's chunks; its answer is not used as a load hit. The save floor starts at the HBM prefix and is lowered to the tier's contiguous hit when the answer arrives (or to 0 if it never does), so chunks the tier evicted are stored again. Read locks are released when the answer arrives; the MP session ends after that, and an unanswered lookup is given up after 120 s, leaving its locks to the read TTL. | | **OFFLOAD_PROFILE** | bool | 0 | Emit `[OFFLOAD-SAVE-PROF]` / `[OFFLOAD-LOAD-PROF]` per-transfer records, and once a minute an `[OFFLOAD-PERF]` line per scheduler with lookup/save/load-skip counters and per-step timings. An empty value reads as off. | | **OFFLOAD_SINGLE_STREAM** | bool | 0 | Experimental: run the staging pack and copy legs on one stream. | diff --git a/tests/test_lmcache_offload_connector.py b/tests/test_lmcache_offload_connector.py index b7124e33db..068b780087 100644 --- a/tests/test_lmcache_offload_connector.py +++ b/tests/test_lmcache_offload_connector.py @@ -7573,54 +7573,37 @@ def test_a_failed_put_discharges_its_reference(self, lmcache_key): codec.put(1, []) assert objs[0].refs == 0 - def test_without_a_local_cpu_backend_a_touch_asks_presence( - self, lmcache_key, caplog + @pytest.mark.parametrize( + "missing", + [ + "local_cpu_backend", + "local_cpu_backend.cpu_lock", + "local_cpu_backend.hot_cache", + "local_cpu_backend.cache_policy", + "cache_policy.update_on_hit", + ], + ) + def test_without_the_lru_internals_a_touch_only_confirms_presence( + self, lmcache_key, missing, caplog ): + """An LMCache whose internals moved keeps serving loads; bind warns + once, naming what is missing, and the lost refresh is counted.""" pool = _LocalCPU() - pool.local_cpu_backend = None - with caplog.at_level(logging.WARNING, logger="atom"): - codec, _ = _lru_codec(pool, max_entries=None) - assert any("no local_cpu_backend" in r.message for r in caplog.records) - pool.hot[codec.key(1)] = _Obj() - assert codec.touch(1) is TouchResult.PRESENT - pool.hot.clear() - assert codec.touch(1) is TouchResult.GONE - - def _assert_degrades(self, pool, missing, caplog): - """Bind warns once naming what is missing; touches then only confirm - presence, and each one is counted rather than passing as a refresh.""" + if missing == "local_cpu_backend": + del pool.local_cpu_backend + elif missing == "cache_policy.update_on_hit": + pool.local_cpu_backend.cache_policy = SimpleNamespace() + else: + setattr(pool.local_cpu_backend, missing.split(".")[1], None) with caplog.at_level(logging.WARNING, logger="atom"): codec, _ = _lru_codec(pool, max_entries=None) bind = [r.message for r in caplog.records if missing in r.message] assert len(bind) == 1 and "state_offload_touches_unrefreshed" in bind[0] assert codec.put(1, []) assert codec.touch(1) is TouchResult.PRESENT - assert codec.touch(1) is TouchResult.PRESENT - assert codec._unrefreshed_touches == 2 - - @pytest.mark.parametrize("dropped", ["cpu_lock", "hot_cache", "cache_policy"]) - def test_a_local_cpu_without_the_lru_internals_degrades_to_presence( - self, lmcache_key, dropped, caplog - ): - """An LMCache whose internals moved keeps serving loads; the lost - refresh shows in the stats instead of taking the engine down.""" - pool = _LocalCPU() - setattr(pool.local_cpu_backend, dropped, None) - self._assert_degrades(pool, f"local_cpu_backend.{dropped}", caplog) - - def test_a_policy_without_update_on_hit_degrades_to_presence( - self, lmcache_key, caplog - ): - pool = _LocalCPU() - pool.local_cpu_backend.cache_policy = SimpleNamespace() - self._assert_degrades(pool, "cache_policy.update_on_hit", caplog) - - def test_a_storage_manager_without_local_cpu_backend_degrades_to_presence( - self, lmcache_key, caplog - ): - pool = _LocalCPU() - del pool.local_cpu_backend - self._assert_degrades(pool, "no local_cpu_backend", caplog) + assert codec._unrefreshed_touches == 1 + pool.hot.clear() + assert codec.touch(1) is TouchResult.GONE def test_the_lru_is_resolved_once_at_bind(self, lmcache_key): pool = _LocalCPU() @@ -7631,31 +7614,24 @@ def test_the_lru_is_resolved_once_at_bind(self, lmcache_key): assert codec.touch(1) is TouchResult.REFRESHED assert pool.evict_oldest() == "kv-chunk" - def test_a_disk_only_image_is_counted_and_warned_once(self, lmcache_key, caplog): - """Present, so still readable, but its LRU was not refreshed: it must - not be silent.""" + def test_only_a_presence_only_touch_is_counted_and_warned_once( + self, lmcache_key, caplog + ): + """A disk-only image is still readable, but its LRU was not refreshed: + that must not be silent, and must not be one warning per touch.""" pool = _LocalCPU() codec, _ = _lru_codec(pool, max_entries=None) - assert codec.put(1, []) - pool.spill_to_disk(codec.key(1)) + assert codec.put(1, []) and codec.put(2, []) + pool.spill_to_disk(codec.key(2)) with caplog.at_level(logging.WARNING, logger="atom"): - assert codec.touch(1) is TouchResult.PRESENT - assert codec.touch(1) is TouchResult.PRESENT + assert codec.touch(1) is TouchResult.REFRESHED + assert codec.touch(2) is TouchResult.PRESENT + assert codec.touch(2) is TouchResult.PRESENT + assert codec.touch(3) is TouchResult.GONE assert codec._unrefreshed_touches == 2 warned = [r for r in caplog.records if "touches_unrefreshed" in r.message] assert len(warned) == 1 - def test_a_refreshed_or_gone_image_is_not_counted(self, lmcache_key, caplog): - pool = _LocalCPU() - codec, _ = _lru_codec(pool, max_entries=None) - assert codec.put(1, []) - with caplog.at_level(logging.WARNING, logger="atom"): - assert codec.touch(1) is TouchResult.REFRESHED - pool.hot.clear() - assert codec.touch(1) is TouchResult.GONE - assert codec._unrefreshed_touches == 0 - assert not [r for r in caplog.records if "touches_unrefreshed" in r.message] - def test_a_disk_only_image_is_stored_back_into_the_cpu_tier(self, lmcache_key): """The put fast path is for an image LocalCPU still holds, not for one only the disk does.""" @@ -7713,18 +7689,9 @@ def test_with_no_tier_every_image_is_gone(self): worker._touch_states(SimpleNamespace(state_touches=[StateTouchId(1, 0)])) assert worker._touch_verdicts == {StateTouchId(1, 0): (False, True)} - def test_a_failing_touch_does_not_erase_the_index(self): - def touch(h): - raise RuntimeError("backend hiccup") - - worker = _touch_worker(SimpleNamespace(codec=SimpleNamespace(touch=touch))) - worker._touch_states(SimpleNamespace(state_touches=[StateTouchId(1, 0)])) - assert worker._touch_verdicts[StateTouchId(1, 0)][0] is True - def test_a_failing_touch_counts_unrefreshed_and_logs_once(self, caplog): - """A drifted LMCache raises on every touch: that must show up in - `touches_unrefreshed`, not as healthy refreshes, and not as one - traceback per touch.""" + """A raising touch keeps the image (resident) but is not a refresh; + a drifted LMCache raises on every touch, so one traceback.""" def touch(h): raise RuntimeError("backend hiccup") @@ -7762,7 +7729,6 @@ def _scheduler(self, monkeypatch): def test_each_touch_is_a_distinct_identity_drained_once(self, monkeypatch): s = self._scheduler(monkeypatch) - assert s.enqueue_state_touches([]) is False assert s.enqueue_state_touches([5, 5]) meta = s.build_connector_meta() assert meta.state_touches == [StateTouchId(5, 0), StateTouchId(5, 1)] @@ -7785,21 +7751,6 @@ def test_one_rank_missing_the_image_makes_it_gone(self, monkeypatch): assert s.take_state_touch_verdicts() == {5: False, 6: True} assert s.take_state_touch_verdicts() == {} - def test_a_presence_only_shard_counts_the_touch_unrefreshed(self, monkeypatch): - from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( - STATE_TOUCH_REFRESH_CHANNEL, - ) - - s = self._scheduler(monkeypatch) - for op, ok in ((StateTouchId(5, 0), True), (StateTouchId(6, 1), False)): - handled = s.connector_completion( - ConnectorCompletion(STATE_TOUCH_REFRESH_CHANNEL, op, ok) - ) - assert handled is None, "handled, but not a terminal save" - assert s.take_state_touches_unrefreshed() == 1 - assert s.take_state_touches_unrefreshed() == 0 - assert s.take_state_touch_verdicts() == {}, "not a residency verdict" - def test_both_touch_reports_reach_quorum_across_tp(self, monkeypatch): """Each rank files a residency and a refresh report per touch; one rank that only found the image present counts the touch unrefreshed @@ -7832,7 +7783,7 @@ def test_both_touch_reports_reach_quorum_across_tp(self, monkeypatch): merged = KVOutputAggregator(world_size=2).aggregate(outputs) s = self._scheduler(monkeypatch) for completion in merged.connector_completions: - s.connector_completion(completion) + assert s.connector_completion(completion) is None, "not a save" assert s.take_state_touch_verdicts() == {5: True} assert s.take_state_touches_unrefreshed() == 1 @@ -7863,7 +7814,3 @@ def test_the_shells_forward_the_unrefreshed_count(self, monkeypatch): MultiConnectorScheduler, "_state_tier_sub", lambda self: None ) assert multi.take_state_touches_unrefreshed() == 0 - - def test_a_negative_round_is_refused(self): - with pytest.raises(ValueError): - StateTouchId(1, -1) diff --git a/tests/test_state_checkpoint.py b/tests/test_state_checkpoint.py index 5bd859d771..b135ddbb69 100644 --- a/tests/test_state_checkpoint.py +++ b/tests/test_state_checkpoint.py @@ -19,8 +19,11 @@ import pytest from conftest import MockConfig -from atom.model_engine import state_offload as so -from atom.model_engine.block_manager import AdmissionProbe, BlockManager +from atom.model_engine.block_manager import ( + AdmissionProbe, + BlockManager, + JointBoundaryDecision, +) from atom.model_engine.block_pool import BlockPool from atom.model_engine.engine_stats import EngineStats from atom.model_engine.page_unit_checkpoint import ( @@ -31,7 +34,7 @@ from atom.model_engine.scheduler import ScheduledBatchOutput, Scheduler from atom.model_engine.sequence import Sequence, SequenceType from atom.model_engine.state_cache import StateCache -from atom.model_engine.state_offload import StateOffloadIndex +from atom.model_engine.state_offload import _TIER_VERIFIED_MAX_S, StateOffloadIndex from atom.model_engine.state_pool import StateSlotPool from atom.model_engine.state_runtime import ( StateRuntime, @@ -4567,12 +4570,12 @@ def test_a_cold_start_without_a_boundary_is_untouched(self): def clock(monkeypatch): """`state_offload`'s monotonic clock, advanced by hand.""" now = [1000.0] - monkeypatch.setattr(so, "monotonic", lambda: now[0]) + monkeypatch.setattr("atom.model_engine.state_offload.monotonic", lambda: now[0]) return now #: Past every confirmation window `tier_verified_for_s` can give. -STALE = so._TIER_VERIFIED_MAX_S + 1 +STALE = _TIER_VERIFIED_MAX_S + 1 def _paged_bm(): @@ -4593,8 +4596,6 @@ class TestAdmissionHold: would load an unconfirmed image, and commits nothing while it does.""" def _joint(self, bm, monkeypatch, *, h=77, from_hbm=False): - from atom.model_engine.block_manager import JointBoundaryDecision - decision = JointBoundaryDecision( boundary_tokens=16, boundary_hash=h, @@ -4668,22 +4669,25 @@ def test_the_fit_probe_path_never_holds(self, clock, monkeypatch): clock[0] += STALE assert bm.record_allocation(_seq(), AdmissionProbe(2, [1, 2])) - def test_an_hbm_hit_holds_only_for_a_tier_sourced_state(self, clock, monkeypatch): - from atom.model_engine.block_manager import JointBoundaryDecision - - bm = _paged_bm() + def _hbm_hit(self, bm, monkeypatch, resident): + """No joint; the HBM hit's state (hash 2) is stale on the tier and in + HBM while `resident` holds it.""" monkeypatch.setattr( bm, "_joint_kv_boundary", lambda *a: JointBoundaryDecision(skip_reason="no_kv_reach"), ) monkeypatch.setattr(bm, "_commit_joint_boundary", lambda seq, d: None) - bm.state_offload.note_stored(2) - clock[0] += STALE - resident = {2} monkeypatch.setattr( bm.paged_state_checkpoints, "contains", lambda h: h in resident ) + bm.state_offload.note_stored(2) + + def test_an_hbm_hit_holds_only_for_a_tier_sourced_state(self, clock, monkeypatch): + bm = _paged_bm() + resident = {2} + self._hbm_hit(bm, monkeypatch, resident) + clock[0] += STALE assert bm.record_allocation( _seq(), AdmissionProbe(2, [1, 2]), defer_unverified=True ) @@ -4692,6 +4696,17 @@ def test_an_hbm_hit_holds_only_for_a_tier_sourced_state(self, clock, monkeypatch _seq(), AdmissionProbe(2, [1, 2]), defer_unverified=True ) + def test_a_replayed_hit_is_not_held(self, clock, monkeypatch): + bm = _paged_bm() + self._hbm_hit(bm, monkeypatch, set()) + monkeypatch.setattr(bm, "state_replay_tokens", 8) + clock[0] += STALE + seq = _seq() + seq.replays_hit = True + assert bm.record_allocation( + seq, AdmissionProbe(2, [1, 2]), defer_unverified=True + ) + def test_the_hold_judges_the_joint_fit_kept(self, clock, monkeypatch): """A joint `_fit_joint` drops is not loaded, so its stale image must not hold an admission that resumes at an HBM-resident hit.""" @@ -4759,26 +4774,6 @@ def test_the_joint_boundary_wins_then_a_replay_then_the_hit(self, monkeypatch): assert bm._state_resume_hash(seq, -1, 2) == -1, "a replay restores none" assert bm._state_resume_hash(seq, 77, 2) == 77 - def test_a_replayed_hit_is_not_held(self, clock, monkeypatch): - from atom.model_engine.block_manager import JointBoundaryDecision - - bm = _paged_bm() - monkeypatch.setattr( - bm, - "_joint_kv_boundary", - lambda *a: JointBoundaryDecision(skip_reason="no_kv_reach"), - ) - monkeypatch.setattr(bm, "_commit_joint_boundary", lambda seq, d: None) - monkeypatch.setattr(bm.paged_state_checkpoints, "contains", lambda h: False) - monkeypatch.setattr(bm, "state_replay_tokens", 8) - bm.state_offload.note_stored(2) - clock[0] += STALE - seq = _seq() - seq.replays_hit = True - assert bm.record_allocation( - seq, AdmissionProbe(2, [1, 2]), defer_unverified=True - ) - def test_the_hold_judges_the_hash_allocate_then_loads(self, clock, monkeypatch): bm = _paged_bm() offload = bm.state_offload diff --git a/tests/test_state_offload_index.py b/tests/test_state_offload_index.py index 71da0df1e8..809d36c4bd 100644 --- a/tests/test_state_offload_index.py +++ b/tests/test_state_offload_index.py @@ -85,13 +85,10 @@ def _index(*hashes, refresh=True, **kw): class TestIndexTouches: - def test_could_serve_queues_nothing(self, clock): - idx = _index(7) - assert idx.could_serve(7) and not idx.could_serve(8) - assert idx.take_touches() == [] - - def test_a_registered_hit_queues_a_touch(self, clock): + def test_only_a_registered_hit_queues_a_touch(self, clock): idx = _index(7) + assert idx.could_serve(7) + assert idx.take_touches() == [], "could_serve is pure" idx.want_touch(7) idx.want_touch(8) # not indexed: nothing to touch assert idx.take_touches() == [7] @@ -99,7 +96,6 @@ def test_a_registered_hit_queues_a_touch(self, clock): def test_no_touches_without_refresh(self, clock): idx = _index(7, refresh=False) - assert idx.could_serve(7) idx.want_touch(7) assert idx.take_touches() == [] @@ -115,49 +111,32 @@ def test_a_hash_voted_every_pass_is_touched_at_an_interval(self, clock): assert idx.take_touches() == [7] assert idx.touches_sent == 2 - def test_a_gone_verdict_drops_the_hash(self, clock): - idx = _index(7) - clock[0] += 1 - idx.want_touch(7) - assert idx.take_touches() == [7] - idx.apply_touch_verdict(7, False) - assert not idx.could_serve(7) - assert idx.touch_misses == 1 - assert idx.stats()["touch_misses"] == 1 - - def test_a_gone_verdict_older_than_a_store_is_ignored(self, clock): - """The touch looked before the store's bytes landed; its verdict must - not erase the image that store just indexed.""" - idx = _index(7) - clock[0] += 1 - idx.want_touch(7) - assert idx.take_touches() == [7] - clock[0] += 0.01 - idx.note_stored(7) - idx.apply_touch_verdict(7, False) - assert idx.could_serve(7) - assert idx.touch_misses == 0 - - def test_a_gone_verdict_for_a_forgotten_then_restored_hash_is_ignored(self, clock): - idx = _index(7) - clock[0] += 1 - idx.want_touch(7) - assert idx.take_touches() == [7] - idx.forget(7) - clock[0] += 0.01 - idx.note_stored(7) - idx.apply_touch_verdict(7, False) - assert idx.could_serve(7) - - def test_a_gone_verdict_after_a_store_and_a_newer_touch_applies(self, clock): + @pytest.mark.parametrize( + "before_touch, after_touch, dropped", + [ + ((), (), True), + ((), ("note_stored",), False), + ((), ("forget", "note_stored"), False), + (("note_stored",), (), True), + ], + ids=["gone", "stored-since", "restored-since", "stored-before"], + ) + def test_a_gone_verdict_drops_the_hash_unless_a_store_landed_since( + self, clock, before_touch, after_touch, dropped + ): idx = _index(7) - clock[0] += 1 - idx.note_stored(7) + for op in before_touch: + clock[0] += 1 + getattr(idx, op)(7) clock[0] += 1 idx.want_touch(7) assert idx.take_touches() == [7] + for op in after_touch: + clock[0] += 0.01 + getattr(idx, op)(7) idx.apply_touch_verdict(7, False) - assert not idx.could_serve(7) + assert idx.could_serve(7) is not dropped + assert idx.stats()["touch_misses"] == int(dropped) def test_a_resident_verdict_refreshes_the_cap_order(self, clock): idx = StateOffloadIndex(max_cpu_entries=19, refresh_hits=True) # cap 3 @@ -173,13 +152,6 @@ def test_a_verdict_for_a_forgotten_hash_is_ignored(self, clock): idx.apply_touch_verdict(7, True) assert not idx.could_serve(7) - def test_stop_refreshing_drops_pending_touches(self, clock): - idx = _index(7) - idx.want_touch(7) - idx.stop_refreshing() - assert idx.take_touches() == [] - assert not idx.needs_verify(7) - class TestNeedsVerify: def test_a_fresh_store_needs_no_verdict(self, clock): @@ -210,75 +182,54 @@ def test_a_gone_verdict_ends_the_hold_by_dropping_the_hash(self, clock): assert not idx.needs_verify(7) and not idx.could_serve(7) def test_a_lost_verdict_holds_no_longer_than_the_lookup_defer(self, clock): - idx = _index(7, defer_s=2) + idx = _index(7) clock[0] += STALE - # A held request re-probes every pass and its touch is re-sent every - # interval; verdicts never come back. Re-sends must not extend the hold. + # Re-probed every pass, re-touched every interval, never answered: + # the re-sends must not extend the hold. held_for, step = 0.0, 0.05 while idx.needs_verify(7): - assert held_for <= 2 + 1e-6, "the hold outlived OFFLOAD_LOOKUP_DEFER_S" - assert idx.could_serve(7) + assert held_for <= idx.defer_s + 1e-6 idx.take_touches() clock[0] += step held_for += step - assert held_for >= 2 + assert held_for >= idx.defer_s assert idx.touches_sent > 1, "the touch was re-sent while held" - assert idx.verify_holds_expired == 1 - assert idx.stats()["verify_holds_expired"] == 1 assert not idx.needs_verify(7), "later admissions are not held again" - assert idx.verify_holds_expired == 1 + assert idx.stats()["verify_holds_expired"] == 1 def test_evidence_after_an_expired_hold_rearms_it(self, clock): - idx = _index(7, defer_s=2) + idx = _index(7) clock[0] += STALE assert idx.needs_verify(7) - clock[0] += 3 + clock[0] += idx.defer_s + 1 assert not idx.needs_verify(7) idx.apply_touch_verdict(7, True) clock[0] += STALE assert idx.needs_verify(7), "a fresh hold, timed from now" - clock[0] += 1 + clock[0] += idx.defer_s / 2 assert idx.needs_verify(7) def test_a_verdict_ends_the_hold_timer(self, clock): - idx = _index(7, defer_s=2) + idx = _index(7) clock[0] += STALE assert idx.needs_verify(7) - clock[0] += 1.5 + clock[0] += idx.defer_s * 0.75 idx.apply_touch_verdict(7, True) clock[0] += STALE assert idx.needs_verify(7) - clock[0] += 1.5 - assert idx.needs_verify(7), "the earlier 1.5 s does not count" - assert idx.verify_holds_expired == 0 + clock[0] += idx.defer_s * 0.75 + assert idx.needs_verify(7), "the earlier wait does not count" - def test_a_zero_defer_never_holds(self, clock, monkeypatch): + def test_the_lookup_defer_is_read_once_at_construction(self, clock, monkeypatch): monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "0") idx = _index(7) - assert idx.defer_s == 0 clock[0] += STALE - assert not idx.needs_verify(7) - - def test_the_lookup_defer_is_read_once_at_construction(self, clock, monkeypatch): - """`needs_verify` is on the admission path; it must not re-read and - re-parse `OFFLOAD_LOOKUP_DEFER_S` on every call.""" + assert not idx.needs_verify(7), "0 disables the hold" monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "2") - idx = _index(7) - assert idx.defer_s == 2.0 - - def unread(): - raise AssertionError("OFFLOAD_LOOKUP_DEFER_S read on the admission path") - - monkeypatch.setattr(so, "verify_defer_s", unread) - monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "0") - clock[0] += STALE - assert idx.needs_verify(7), "the value read at construction still holds" + assert not idx.needs_verify(7), "not re-read on the admission path" def test_an_index_without_refresh_never_reads_the_defer(self, monkeypatch): - def unread(): - raise AssertionError("read for an index that never holds") - - monkeypatch.setattr(so, "verify_defer_s", unread) + monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "not-a-number") assert StateOffloadIndex().defer_s == 0 def test_a_completed_load_counts_as_verified(self, clock): @@ -288,40 +239,33 @@ def test_a_completed_load_counts_as_verified(self, clock): idx.complete_load("r") assert not idx.needs_verify(7) + def test_a_hold_is_counted(self, clock): + idx = _index(7) + assert not idx.hold_admission(-1), "nothing to load" + assert not idx.hold_admission(7), "confirmed by its store" + clock[0] += STALE + assert idx.hold_admission(7) + assert idx.stats()["loads_deferred_for_verify"] == 1 -class TestVerifiedWindow: - """How long a confirmed image is trusted scales with the LMCache pool it - has to age out of, within [5 s, 60 s].""" - - GIB = 1 << 30 - - def test_the_agentx_pool_keeps_about_the_measured_window(self): - assert so.tier_verified_for_s(int(208.6 * self.GIB)) == pytest.approx(52.15) - def test_the_window_grows_with_the_pool_between_its_bounds(self): - sizes = [8, 20, 32, 64, 100, 208.6, 240, 512, 2048] - windows = [so.tier_verified_for_s(int(g * self.GIB)) for g in sizes] - assert windows == sorted(windows) - assert windows[0] == windows[1] == so._TIER_VERIFIED_MIN_S - assert windows[2] == pytest.approx(8.0) - assert windows[-3] == windows[-1] == so._TIER_VERIFIED_MAX_S +@pytest.mark.parametrize( + "pool_gib, window", [(0, 5.0), (8, 5.0), (32, 8.0), (208.6, 52.15), (512, 60.0)] +) +def test_the_verified_window_scales_with_the_pool_within_bounds(pool_gib, window): + assert so.tier_verified_for_s(pool_gib) == pytest.approx(window) - def test_an_unknown_pool_gets_the_floor(self): - assert so.tier_verified_for_s(0) == so._TIER_VERIFIED_MIN_S - def test_the_index_holds_by_its_own_window(self, clock): - idx = _index(7, verified_for_s=so.tier_verified_for_s(32 * self.GIB), defer_s=2) - assert idx.verified_for_s == pytest.approx(8.0) - clock[0] += 7.5 - assert not idx.needs_verify(7) - clock[0] += 1 - assert idx.needs_verify(7) +def test_the_index_holds_by_its_own_window(clock): + idx = _index(7, verified_for_s=8.0) + clock[0] += 7.5 + assert not idx.needs_verify(7) + clock[0] += 1 + assert idx.needs_verify(7) class TestTouchPlumbing: - """The index's two calls on the connector, made by the scheduler once a - pass (`publish_touches` before the metadata is built, - `collect_touch_verdicts` before the next admission).""" + """The scheduler's two calls a pass: `publish_touches` before the metadata + is built, `collect_touch_verdicts` before the next admission.""" def test_voted_hashes_reach_the_connector(self, clock): idx = _index(7) @@ -336,67 +280,40 @@ def test_a_connector_without_touches_turns_the_refresh_off(self, clock): idx = _index(7) idx.want_touch(7) idx.publish_touches(SimpleNamespace()) - assert not idx.refresh_hits + idx.want_touch(7) + assert idx.take_touches() == [] clock[0] += STALE assert not idx.needs_verify(7), "nothing would ever answer the hold" - def test_verdicts_are_applied_to_the_index(self, clock): + def test_verdicts_and_presence_only_touches_are_collected(self, clock): idx = _index(7, 8) clock[0] += 1 idx.want_touch(7) idx.want_touch(8) assert sorted(idx.take_touches()) == [7, 8] - idx.collect_touch_verdicts( - SimpleNamespace(take_state_touch_verdicts=lambda: {7: False, 8: True}) - ) - assert not idx.could_serve(7) and idx.could_serve(8) - - def test_presence_only_touches_reach_the_stats_line(self, clock): - idx = _index(7) connector = SimpleNamespace( - take_state_touch_verdicts=lambda: {7: True}, + take_state_touch_verdicts=lambda: {7: False, 8: True}, take_state_touches_unrefreshed=lambda: 3, ) - assert idx.stats()["touches_unrefreshed"] == 0 idx.collect_touch_verdicts(connector) idx.collect_touch_verdicts(connector) + assert not idx.could_serve(7) and idx.could_serve(8) assert idx.stats()["touches_unrefreshed"] == 6 - def test_an_offered_load_registers_a_touch(self, clock): - idx = _index(7) - assert idx.request_load("r", 7) - assert idx.take_touches() == [7] - - def test_a_hold_is_counted(self, clock): - idx = _index(7) - assert not idx.hold_admission(-1), "nothing to load" - assert not idx.hold_admission(7), "confirmed by its store" - clock[0] += STALE - assert idx.hold_admission(7) - assert idx.loads_deferred_for_verify == 1 - assert idx.stats()["loads_deferred_for_verify"] == 1 - - -# ── the startup budget check ─────────────────────────────────────────────── - GIB = 1 << 30 -class TestBudgetWarning: - def test_a_cap_below_half_the_pool_is_fine(self): - assert state_tier_budget_warning(59 * GIB, 208 * GIB, 112 << 20) is None - - def test_a_cap_the_size_of_the_pool_bounds_nothing(self): - """#2562's accuracy run: a 32 GiB pool with the 32 GiB default cap.""" - msg = state_tier_budget_warning(32 * GIB, 32 * GIB, 112 << 20) - assert "bounds nothing" in msg - assert "LMCACHE_MAX_LOCAL_CPU_SIZE=32.0" in msg - - def test_a_cap_over_half_the_pool_is_called_out(self): - msg = state_tier_budget_warning(40 * GIB, 64 * GIB, 112 << 20) - assert msg is not None and "more than half" in msg - - def test_unknown_sizes_are_not_judged(self): - assert state_tier_budget_warning(32 * GIB, 0, 1) is None - assert state_tier_budget_warning(0, 32 * GIB, 1) is None +@pytest.mark.parametrize( + "state, pool, says", + [ + (59, 208, None), + (32, 32, "bounds nothing"), # #2562's accuracy run: the default cap + (40, 64, "more than half"), + (32, 0, None), + (0, 32, None), + ], +) +def test_the_state_cap_is_judged_against_the_pool(state, pool, says): + msg = state_tier_budget_warning(state * GIB, pool * GIB) + assert msg is None if says is None else says in msg From 024e3e8c32850e223164b481df458bbd5ab69135 Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Mon, 2 Nov 2026 05:06:14 +0000 Subject: [PATCH 13/18] state offload: count unrefreshed touches in the connector statistics Per the independent review of the cleanup: - The presence-only touch count rides KimiK3OffloadScheduler.get_statistics() instead of its own take method forwarded by the shell and multi. - An offered load no longer queues a touch: the load is its own verdict and the codec's get refreshes the LRU. - needs_verify folds into hold_admission, its only caller; refresh_hits and verified_for_s become one verified_for_s argument; the replay check reuses tier_resume_replays. - The codec warns once per cause; StateTouchId.round is renamed generation; the touch methods join StateOffloadFace; the budget warning takes GiB. - Tests share the offload clock in conftest and the K3 scheduler helper, and cover the touch state leaving with its hash. --- atom/kv_transfer/disaggregation/aggregator.py | 2 +- .../disaggregation/multi/multi_connector.py | 7 - atom/kv_transfer/disaggregation/types.py | 5 +- atom/kv_transfer/offload/README.md | 6 +- atom/kv_transfer/offload/_offload_common.py | 6 +- atom/kv_transfer/offload/connector.py | 5 - .../offload/hybrid/kimi_k3/connector.py | 35 +-- .../offload/hybrid/kimi_k3/state_object.py | 30 +-- atom/kv_transfer/offload/metadata.py | 3 +- atom/model_engine/block_manager.py | 7 +- atom/model_engine/scheduler.py | 19 +- atom/model_engine/state_offload.py | 88 +++----- tests/conftest.py | 13 ++ tests/test_lmcache_offload_connector.py | 95 +++----- tests/test_state_checkpoint.py | 77 +++---- tests/test_state_offload_index.py | 213 +++++++++--------- 16 files changed, 268 insertions(+), 343 deletions(-) diff --git a/atom/kv_transfer/disaggregation/aggregator.py b/atom/kv_transfer/disaggregation/aggregator.py index b614ba030e..d4fcaaf58f 100644 --- a/atom/kv_transfer/disaggregation/aggregator.py +++ b/atom/kv_transfer/disaggregation/aggregator.py @@ -157,7 +157,7 @@ def __init__( terminal_tombstone_limit, lambda key: isinstance(key, LoadOperationId), ) - # Touch ids are never reused (fresh round), and tombstoning them would + # 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, diff --git a/atom/kv_transfer/disaggregation/multi/multi_connector.py b/atom/kv_transfer/disaggregation/multi/multi_connector.py index 0d4a6f6ae7..c5b2c5840d 100644 --- a/atom/kv_transfer/disaggregation/multi/multi_connector.py +++ b/atom/kv_transfer/disaggregation/multi/multi_connector.py @@ -670,13 +670,6 @@ def take_state_touch_verdicts(self) -> dict: 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. diff --git a/atom/kv_transfer/disaggregation/types.py b/atom/kv_transfer/disaggregation/types.py index 177041a521..a842f3fb45 100644 --- a/atom/kv_transfer/disaggregation/types.py +++ b/atom/kv_transfer/disaggregation/types.py @@ -128,11 +128,12 @@ def __post_init__(self) -> None: 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. `round` keeps repeated touches of one hash distinct identities. + shard is. `generation` keeps repeated touches of one hash distinct + identities. """ prefix_hash: int - round: int + generation: int ConnectorCompletionId = ( diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index fd534bf0b8..2c7492e7ea 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -131,9 +131,9 @@ 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 -`state_offload_touches_unrefreshed`. `OFFLOAD_STATE_CPU_SIZE` caps the images' -share of the pool and reserves nothing; the worker warns when it is not below -half of the pool. +the connector statistic `state_touches_unrefreshed`. `OFFLOAD_STATE_CPU_SIZE` +caps the images' share of the pool and reserves nothing; the worker warns when it +is more than half of the pool. ## LMCache Multiprocess (`lmcache_mp`) diff --git a/atom/kv_transfer/offload/_offload_common.py b/atom/kv_transfer/offload/_offload_common.py index cd26767a23..a5e70ba0ca 100644 --- a/atom/kv_transfer/offload/_offload_common.py +++ b/atom/kv_transfer/offload/_offload_common.py @@ -426,7 +426,7 @@ class StateOffloadFace(ABC): `QwenOffloadScheduler` inherits that, so `isinstance(impl, StateOffloadFace)` is the honest predicate. - Narrow on purpose -- only the four tier methods. Other impl attributes the + Narrow on purpose -- only the tier methods. Other impl attributes the scheduler reads (e.g. `max_pending_saves`) stay plain shell forwards rather than joining this face, which is exclusively the KDA state-tier contract. """ @@ -439,6 +439,10 @@ def enqueue_state_stores(self, stores) -> bool: ... def take_state_reports(self) -> tuple[set[int], set[int]]: ... @abstractmethod def take_state_source_releases(self) -> set: ... + @abstractmethod + def enqueue_state_touches(self, hashes) -> bool: ... + @abstractmethod + def take_state_touch_verdicts(self) -> dict[int, bool]: ... class OffloadSchedulerMixin(ABC): diff --git a/atom/kv_transfer/offload/connector.py b/atom/kv_transfer/offload/connector.py index dacf39499a..065a75a963 100644 --- a/atom/kv_transfer/offload/connector.py +++ b/atom/kv_transfer/offload/connector.py @@ -345,11 +345,6 @@ def take_state_touch_verdicts(self) -> dict: 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.""" - 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 diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py index 049fb737d1..b839c5716a 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py @@ -132,8 +132,8 @@ def __init__(self, config) -> None: self._store_failed_no_tier: set[StateStoreOperationId] = set() # Touch verdicts since the last `get_finished`: (resident, refreshed). self._touch_verdicts: dict[StateTouchId, tuple[bool, bool]] = {} - # Touches whose codec call raised; only the first is logged. - self._touch_failures = 0 + # 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. @@ -264,9 +264,7 @@ def _build_state_tier(self, transfer_tensors) -> None: # checkpoints cannot consume nearly the whole paged-KV CPU pool. codec.bind_storage_manager(self._engine.storage_manager) pool_gib = float(getattr(cfg, "max_local_cpu_size", 0) or 0) - budget_warning = state_tier_budget_warning( - state_bytes, int(pool_gib * (1 << 30)) - ) + budget_warning = state_tier_budget_warning(state_bytes / (1 << 30), pool_gib) if budget_warning: logger.warning("kimi_k3 offload: %s", budget_warning) # No index here: StateOffloadIndex lives in the engine process; both @@ -329,12 +327,12 @@ def _touch_states(self, metadata) -> None: except Exception: # Says nothing about the bytes: kept, not refreshed. One # traceback, since a drifted LMCache raises on every touch. - self._touch_failures += 1 - if self._touch_failures == 1: + 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_offload_touches_unrefreshed", + "state_touches_unrefreshed", touch_id.prefix_hash, exc_info=True, ) @@ -676,7 +674,7 @@ def __init__(self, config) -> None: self._pending_state_loads: list[tuple] = [] self._pending_state_stores: list[tuple] = [] self._pending_state_touches: list[StateTouchId] = [] - self._touch_round = 0 + 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 @@ -717,10 +715,12 @@ def enqueue_state_stores(self, stores) -> bool: return True def enqueue_state_touches(self, hashes) -> bool: - """Queue touches for the next metadata, each under its own round.""" + """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_round)) - self._touch_round += 1 + 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]: @@ -729,11 +729,12 @@ def take_state_touch_verdicts(self) -> dict[int, bool]: self._state_touch_verdicts = {} return verdicts - def take_state_touches_unrefreshed(self) -> int: - """Touches since the last call that only confirmed presence.""" - n = self._state_touches_unrefreshed - self._state_touches_unrefreshed = 0 - return n + 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() diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py index 90ce8e303a..83a2e8fbdb 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py @@ -54,7 +54,7 @@ def _local_cpu_lru(storage_manager): logger.warning( "state offload: LMCache's storage manager has no %s, so state touches " "only confirm presence and refresh no eviction order (every touch is " - "counted in state_offload_touches_unrefreshed)", + "counted in state_touches_unrefreshed)", ", ".join(missing), ) return None @@ -96,7 +96,7 @@ def __init__( self._worker_id = int(worker_id) self._layout_id = layout_id self._misfit_reads = 0 - self._unrefreshed_touches = 0 + self._warned_unrefreshed = False self._storage = None self._local_cpu = None self._max_entries = ( @@ -191,23 +191,23 @@ def touch(self, h: int) -> TouchResult: """Refresh `h`'s image in LMCache's LRU, as a KV lookup does for its chunks; falsy (`GONE`) when it is gone. A state hit is answered by `StateOffloadIndex`, so otherwise LMCache never sees the image used. - - `update_on_hit` under `cpu_lock`, not `contains(pin=True)` + - `touch_cache()`: a pin needs an unpin, and `keys_in_request` belongs - to the concurrent KV lookup. """ if self._storage is None: return TouchResult.GONE result = self._touch_key(self.key(h)) - if result is TouchResult.PRESENT: - self._unrefreshed_touches += 1 - if self._unrefreshed_touches == 1: - logger.warning( - "state offload: a touched state image is outside " - "LocalCPU's hot cache, so the touch only confirmed it is " - "present and did not refresh its eviction order (counted " - "as state_offload_touches_unrefreshed)" - ) + # Without LocalCPU's internals `_local_cpu_lru` already warned. + if ( + result is TouchResult.PRESENT + and self._local_cpu is not None + and not self._warned_unrefreshed + ): + self._warned_unrefreshed = True + logger.warning( + "state offload: a touched state image is outside " + "LocalCPU's hot cache, so the touch only confirmed it is " + "present and did not refresh its eviction order (counted " + "as state_touches_unrefreshed)" + ) return result def _touch_key(self, key) -> TouchResult: diff --git a/atom/kv_transfer/offload/metadata.py b/atom/kv_transfer/offload/metadata.py index 5cfab3aaf2..2eedb64097 100644 --- a/atom/kv_transfer/offload/metadata.py +++ b/atom/kv_transfer/offload/metadata.py @@ -21,6 +21,7 @@ LoadOperationId, ReqId, SaveOperationId, + StateTouchId, ) @@ -218,7 +219,7 @@ def __init__(self) -> None: self.state_stores: list[tuple] = [] # StateTouchId per voted state image, refreshed and reported by each # worker. - self.state_touches: list = [] + self.state_touches: list[StateTouchId] = [] def add_request(self, meta: LMCacheReqMeta) -> None: self.requests.append(meta) diff --git a/atom/model_engine/block_manager.py b/atom/model_engine/block_manager.py index 0d3ecd84ad..4f8c2d168e 100644 --- a/atom/model_engine/block_manager.py +++ b/atom/model_engine/block_manager.py @@ -516,7 +516,6 @@ def __init__( if checkpoint_spec is not None else None ), - refresh_hits=True, verified_for_s=state_tier_verified_for_s(config), ) # Attached rather than passed at construction: the coordinator is @@ -1613,11 +1612,7 @@ def _state_resume_hash(self, seq: Sequence, joint_hash: int, hit_hash: int) -> i """ if joint_hash != -1: return joint_hash - if self.state.transfer.replays or ( - self.state_replay_tokens and seq.replays_hit - ): - return -1 - return hit_hash + return -1 if self.tier_resume_replays(seq) else hit_hash def retreat_joint_to_hbm(self, seq: Sequence) -> bool: """Drop a failed unaligned joint and resume at the HBM hit it lowered. diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index 9159154ff7..30f62c6b5a 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -3779,8 +3779,8 @@ def _publish_state_stores(self) -> None: self.block_manager.settle_state_store(op, ok=False, attempted=False) def _publish_state_loads(self) -> None: - """Hand this pass's state-tier loads to the connector, before it builds - its metadata. + """Hand this pass's state-tier loads, and the touches the admission + pass registered, to the connector before it builds its metadata. Drained here rather than at the park, because the connector publishes once per pass and a load handed over twice would write the same entry @@ -3791,10 +3791,11 @@ def _publish_state_loads(self) -> None: # Guarded like the `_settle_state_load` / `_publish_state_stores` # siblings: a scheduler built without a block manager (the connector # test doubles) would otherwise AttributeError on the bare deref. - offload = getattr(getattr(self, "block_manager", None), "state_offload", None) + bm = getattr(self, "block_manager", None) + offload = getattr(bm, "state_offload", None) if offload is not None: offload.publish_touches(self.kv_connector) - take = getattr(getattr(self, "block_manager", None), "take_state_loads", None) + take = getattr(bm, "take_state_loads", None) if take is None: return loads = take() @@ -4906,6 +4907,11 @@ def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput): # offload leg. self._drain_source_safe_releases() + bm = getattr(self, "block_manager", None) + offload = getattr(bm, "state_offload", None) + if offload is not None: + # Touch verdicts, before the next admission pass reads them. + offload.collect_touch_verdicts(self.kv_connector) # The state-offload store reports, not keyed by request: by the time a # store lands its owner is long gone and only the hash remains. They # ride the connector's completion channels, so they are drained from the @@ -4914,11 +4920,6 @@ def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput): # Indexed only on the report, never at submission: a hash advertised # before its bytes exist parks the next request over that prefix # against a `get` that must miss. - bm = getattr(self, "block_manager", None) - offload = getattr(bm, "state_offload", None) - if offload is not None: - # Before the next admission pass reads them. - offload.collect_touch_verdicts(self.kv_connector) take = getattr(self.kv_connector, "take_state_reports", None) if offload is not None and take is not None: indexed, failed = take() diff --git a/atom/model_engine/state_offload.py b/atom/model_engine/state_offload.py index 2e1c405639..39b98f8c92 100644 --- a/atom/model_engine/state_offload.py +++ b/atom/model_engine/state_offload.py @@ -24,6 +24,7 @@ #: Seconds a confirmed image is trusted per GiB of LMCache pool: a quarter of #: the pool's turnover at 1 GiB/s, four times the fill rate AgentX c64 measured. _TIER_VERIFIED_S_PER_GIB = 0.25 +#: Well above a touch round trip / the largest window measured (AgentX c64). _TIER_VERIFIED_MIN_S = 5.0 _TIER_VERIFIED_MAX_S = 60.0 @@ -38,11 +39,11 @@ class StateOffloadIndex: can drop bytes at any time, so on its own a hash here means "was stored once", never "is still there". - With `refresh_hits`, every voted hit is also touched: the worker refreshes - the image in LMCache's LRU, as a KV lookup does for its chunks, and reports - whether it was still there (`apply_touch_verdict`). A false positive that - slips through still costs one park/unpark and is handled by the - `failed_loading` path (`fail_load` -> `forget`). + With a `verified_for_s` window, every voted hit is also touched: the worker + refreshes the image in LMCache's LRU, as a KV lookup does for its chunks, + and reports whether it was still there (`apply_touch_verdict`). A false + positive that slips through still costs one park/unpark and is handled by + the `failed_loading` path (`fail_load` -> `forget`). """ def __init__( @@ -51,8 +52,7 @@ def __init__( can_store: bool = True, can_load: bool = True, max_cpu_entries: int | None = None, - refresh_hits: bool = False, - verified_for_s: float = _TIER_VERIFIED_MIN_S, + verified_for_s: float | None = None, ) -> None: # The two legs are separately granted. A `kv_producer` role saves and # never loads, a `kv_consumer` loads and never saves, and offering a @@ -105,9 +105,9 @@ def __init__( # on live transfers and is set too low. self.stores_untrusted = 0 # Off for indexes whose connector carries no touches (the vLLM plugin). - self.refresh_hits = bool(refresh_hits) - self.verified_for_s = float(verified_for_s) - # Read once: `needs_verify` runs on the admission path. Not read + self.refresh_hits = verified_for_s is not None + self.verified_for_s = float(verified_for_s or 0.0) + # Read once: `hold_admission` runs on the admission path. Not read # without the refresh, which never holds. self.defer_s = float(envs.OFFLOAD_LOOKUP_DEFER_S) if self.refresh_hits else 0.0 self._touch_wanted: set[int] = set() @@ -122,7 +122,6 @@ def __init__( self._hold_expired: set[int] = set() self.touches_sent = 0 self.touch_misses = 0 - self.touches_unrefreshed = 0 self.loads_deferred_for_verify = 0 self.verify_holds_expired = 0 @@ -160,7 +159,8 @@ def _drop_touch_state(self, h: int) -> None: self._end_hold(h) def _mark_verified(self, h: int) -> None: - self._verified_at[h] = monotonic() + if self.refresh_hits: + self._verified_at[h] = monotonic() self._end_hold(h) def _end_hold(self, h: int) -> None: @@ -187,8 +187,8 @@ def could_serve(self, h: int) -> bool: # ------------------------------- touches ------------------------------- # def want_touch(self, h: int) -> None: - """A voter counted `h` as a hit (the resumable scan, or `request_load`); - under `refresh_hits`, queue it for `take_touches`.""" + """The resumable scan counted `h` as a hit; under `refresh_hits`, queue + it for `take_touches`.""" if self.refresh_hits and h in self.hashes: self._touch_wanted.add(h) @@ -244,10 +244,6 @@ def publish_touches(self, connector) -> None: type(connector).__name__, ) self.refresh_hits = False - self._touch_wanted.clear() - self._touch_sent_at.clear() - self._hold_started_at.clear() - self._hold_expired.clear() def collect_touch_verdicts(self, connector) -> None: """Apply the verdicts `connector` reached TP quorum on.""" @@ -256,15 +252,12 @@ def collect_touch_verdicts(self, connector) -> None: return for h, resident in take().items(): self.apply_touch_verdict(h, resident) - unrefreshed = getattr(connector, "take_state_touches_unrefreshed", None) - if unrefreshed is not None: - self.touches_unrefreshed += int(unrefreshed()) - def needs_verify(self, h: int) -> bool: - """Whether an admission about to load `h` waits a pass for its touch - verdict: `h` not confirmed within `verified_for_s`. If the verdict is - gone, `h` has left the index and the retry plans a lower boundary - instead of parking on a miss. + def hold_admission(self, h: int) -> bool: + """Whether an admission restoring its state from `h` (-1: nothing + outside HBM) waits a pass for its touch verdict: `h` not confirmed + within `verified_for_s`. If the verdict is gone, `h` has left the index + and the retry plans a lower boundary instead of parking on a miss. Bounded by `defer_s` from the first hold, not from the latest touch (re-sent every `_TIER_TOUCH_INTERVAL_S`, so it would never run out); @@ -295,13 +288,6 @@ def needs_verify(self, h: int) -> bool: return False # Ask for a touch; `take_touches` throttles the resend. self._touch_wanted.add(h) - return True - - def hold_admission(self, h: int) -> bool: - """`needs_verify` for an admission restoring its state from `h` (-1: - nothing outside HBM), counted.""" - if h == -1 or not self.needs_verify(h): - return False self.loads_deferred_for_verify += 1 return True @@ -311,8 +297,7 @@ def request_load(self, req_id, h: int) -> bool: The guard between believing and delivering: a load is resolved only by a worker report, so offering one for a hash never stored would park the - request against bytes no `get` can produce. An offered load is a hit - (`want_touch`). + request against bytes no `get` can produce. """ if not self.could_serve(h): # A store-only role (can_load False) or a hash never stored. Voting @@ -332,7 +317,6 @@ def request_load(self, req_id, h: int) -> bool: return False self.pending_loads[req_id] = int(h) self.loads_attempted += 1 - self.want_touch(h) return True def complete_load(self, req_id) -> None: @@ -400,27 +384,20 @@ def stats(self) -> dict[str, int]: # far past `indexed` means the state budget is smaller than the # reusable working set. "hashes_evicted": self.hashes_evicted, - # Hit refresh. `touch_misses`: images LMCache dropped that left the - # index before anything parked on them (`loads_failed` is what - # still got past). `touches_unrefreshed`: confirmed present but - # not refreshed. `verify_holds_expired`: holds that got no verdict - # within `defer_s`; non-zero means verdicts are being lost. + # Hit refresh. `touch_misses`: images a touch found gone, which + # left the index (one a load was already offered on also counts in + # `loads_failed`). `verify_holds_expired`: holds that got no + # verdict within `defer_s`; non-zero means verdicts are being lost. "touches_sent": self.touches_sent, "touch_misses": self.touch_misses, - "touches_unrefreshed": self.touches_unrefreshed, "loads_deferred_for_verify": self.loads_deferred_for_verify, "verify_holds_expired": self.verify_holds_expired, } def tier_verified_for_s(pool_gib: float) -> float: - """How long an image counts as resident after the last evidence (store, - load, touch), which left it at the young end of LMCache's LRU. - - A share of the pool's turnover, clamped: the floor stays well above a - touch round trip, the ceiling is the largest window measured (AgentX c48 / - c64). Either error costs a held pass or a missed load, not wrong output. - """ + """How long an image counts as resident after the last store, load or + touch; either error costs a held pass or a missed load, not wrong output.""" window = pool_gib * _TIER_VERIFIED_S_PER_GIB return min(max(window, _TIER_VERIFIED_MIN_S), _TIER_VERIFIED_MAX_S) @@ -599,24 +576,23 @@ def state_tier_cpu_bytes() -> int: return int(float(os.environ.get("OFFLOAD_STATE_CPU_SIZE", "32")) * (1 << 30)) -def state_tier_budget_warning(state_bytes: int, pool_bytes: int) -> str | None: +def state_tier_budget_warning(state_gib: float, pool_gib: float) -> str | None: """Why the state cap and the LMCache pool it shares do not fit together (e.g. `LMCACHE_MAX_LOCAL_CPU_SIZE=32` with the 32 GiB default cap), or None.""" - if pool_bytes <= 0 or state_bytes <= 0: + if pool_gib <= 0: return None - gib = float(1 << 30) sizes = ( - f"OFFLOAD_STATE_CPU_SIZE={state_bytes / gib:.1f} GiB, " - f"LMCACHE_MAX_LOCAL_CPU_SIZE={pool_bytes / gib:.1f} GiB" + f"OFFLOAD_STATE_CPU_SIZE={state_gib:.1f} GiB, " + f"LMCACHE_MAX_LOCAL_CPU_SIZE={pool_gib:.1f} GiB" ) - if state_bytes >= pool_bytes: + if state_gib >= pool_gib: return ( f"{sizes}: the state cap is not below the pool, so it bounds " "nothing and state images can take the whole pool from paged KV. " "Lower OFFLOAD_STATE_CPU_SIZE or raise LMCACHE_MAX_LOCAL_CPU_SIZE." ) - if 2 * state_bytes > pool_bytes: + if 2 * state_gib > pool_gib: return ( f"{sizes}: state images may take more than half of the pool, " "leaving paged KV the smaller share under pressure." diff --git a/tests/conftest.py b/tests/conftest.py index 209a2a47f6..22ca235a7d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -31,6 +31,7 @@ from atom.model_engine.block_manager import BlockManager from atom.model_engine.scheduler import Scheduler from atom.model_engine.sequence import Sequence +from atom.model_engine.state_offload import _TIER_VERIFIED_MAX_S from atom.sampling_params import SamplingParams from atom.utils.region_graph import ReqGraph @@ -247,6 +248,18 @@ def keep_envs_lazy(): delattr(envs, name) +#: Past every confirmation window `tier_verified_for_s` can give. +OFFLOAD_STALE_S = _TIER_VERIFIED_MAX_S + 1 + + +@pytest.fixture +def offload_clock(monkeypatch): + """`state_offload`'s monotonic clock, advanced by hand.""" + now = [1000.0] + monkeypatch.setattr("atom.model_engine.state_offload.monotonic", lambda: now[0]) + return now + + @pytest.fixture def seq_factory(): """Factory for creating Sequence objects with sensible defaults.""" diff --git a/tests/test_lmcache_offload_connector.py b/tests/test_lmcache_offload_connector.py index 068b780087..4cfddc9011 100644 --- a/tests/test_lmcache_offload_connector.py +++ b/tests/test_lmcache_offload_connector.py @@ -5825,8 +5825,9 @@ def _k3_scheduler() -> KimiK3OffloadScheduler: s._pending_state_loads = [] s._pending_state_stores = [] s._pending_state_touches = [] - s._touch_round = 0 + s._touch_generation = 0 s._state_touch_verdicts = {} + s._state_touches_unrefreshed = 0 s._save_inflight_since = {} s._save_stalled = False s._warned_save_stalled = False @@ -6161,6 +6162,13 @@ def test_no_more_saves_go_out_than_the_pool_can_afford_to_pin(): # ── the state key: build safety and namespace separation ────────────────── +class _LRUPolicy: + """LMCache's `LRUCachePolicy.update_on_hit`.""" + + def update_on_hit(self, key, cache) -> None: + cache.move_to_end(key) + + class _FakeStorage: """The `storage_manager` calls the codec makes, over a plain dict, with the LocalCPU LRU internals `bind_storage_manager` requires.""" @@ -6168,11 +6176,7 @@ class _FakeStorage: def __init__(self) -> None: self.objects: OrderedDict = OrderedDict() self.local_cpu_backend = SimpleNamespace( - cpu_lock=threading.Lock(), - hot_cache=self.objects, - cache_policy=SimpleNamespace( - update_on_hit=lambda key, cache: cache.move_to_end(key) - ), + cpu_lock=threading.Lock(), hot_cache=self.objects, cache_policy=_LRUPolicy() ) def batched_put(self, keys, objs) -> None: @@ -7371,13 +7375,6 @@ def get_size(self) -> int: return self.size -class _LRUPolicy: - """LMCache's `LRUCachePolicy.update_on_hit`.""" - - def update_on_hit(self, key, cache) -> None: - cache.move_to_end(key) - - class _LocalCPU: """`storage_manager` over LocalCPU's rules. @@ -7587,7 +7584,7 @@ def test_without_the_lru_internals_a_touch_only_confirms_presence( self, lmcache_key, missing, caplog ): """An LMCache whose internals moved keeps serving loads; bind warns - once, naming what is missing, and the lost refresh is counted.""" + once, naming what is missing, and a touch adds no second warning.""" pool = _LocalCPU() if missing == "local_cpu_backend": del pool.local_cpu_backend @@ -7597,11 +7594,10 @@ def test_without_the_lru_internals_a_touch_only_confirms_presence( setattr(pool.local_cpu_backend, missing.split(".")[1], None) with caplog.at_level(logging.WARNING, logger="atom"): codec, _ = _lru_codec(pool, max_entries=None) - bind = [r.message for r in caplog.records if missing in r.message] - assert len(bind) == 1 and "state_offload_touches_unrefreshed" in bind[0] - assert codec.put(1, []) - assert codec.touch(1) is TouchResult.PRESENT - assert codec._unrefreshed_touches == 1 + assert codec.put(1, []) + assert codec.touch(1) is TouchResult.PRESENT + warned = [r.message for r in caplog.records] + assert len(warned) == 1 and missing in warned[0] pool.hot.clear() assert codec.touch(1) is TouchResult.GONE @@ -7614,9 +7610,7 @@ def test_the_lru_is_resolved_once_at_bind(self, lmcache_key): assert codec.touch(1) is TouchResult.REFRESHED assert pool.evict_oldest() == "kv-chunk" - def test_only_a_presence_only_touch_is_counted_and_warned_once( - self, lmcache_key, caplog - ): + def test_a_presence_only_touch_is_warned_once(self, lmcache_key, caplog): """A disk-only image is still readable, but its LRU was not refreshed: that must not be silent, and must not be one warning per touch.""" pool = _LocalCPU() @@ -7628,7 +7622,6 @@ def test_only_a_presence_only_touch_is_counted_and_warned_once( assert codec.touch(2) is TouchResult.PRESENT assert codec.touch(2) is TouchResult.PRESENT assert codec.touch(3) is TouchResult.GONE - assert codec._unrefreshed_touches == 2 warned = [r for r in caplog.records if "touches_unrefreshed" in r.message] assert len(warned) == 1 @@ -7644,7 +7637,6 @@ def test_a_disk_only_image_is_stored_back_into_the_cpu_tier(self, lmcache_key): assert pool.puts == 2 and len(objs) == 2 assert codec.key(1) in pool.hot assert released == [True] - assert codec._unrefreshed_touches == 0, "a store is not a touch verdict" # ── the connector: touches out in the metadata, verdicts back by quorum ──── @@ -7658,7 +7650,7 @@ def _touch_worker(tier): c = KimiK3OffloadConnector.__new__(KimiK3OffloadConnector) c._state_tier = tier c._touch_verdicts = {} - c._touch_failures = 0 + c._warned_touch_failure = False return c @@ -7689,7 +7681,7 @@ def test_with_no_tier_every_image_is_gone(self): worker._touch_states(SimpleNamespace(state_touches=[StateTouchId(1, 0)])) assert worker._touch_verdicts == {StateTouchId(1, 0): (False, True)} - def test_a_failing_touch_counts_unrefreshed_and_logs_once(self, caplog): + def test_a_failing_touch_is_unrefreshed_and_logged_once(self, caplog): """A raising touch keeps the image (resident) but is not a refresh; a drifted LMCache raises on every touch, so one traceback.""" @@ -7702,28 +7694,15 @@ def touch(h): worker._touch_states(SimpleNamespace(state_touches=ids[:2])) worker._touch_states(SimpleNamespace(state_touches=ids[2:])) assert worker._touch_verdicts == {i: (True, False) for i in ids} - assert worker._touch_failures == 3 assert [r.exc_info is not None for r in caplog.records] == [True] def _scheduler(self, monkeypatch): - from atom.kv_transfer.offload.dense.connector import DenseOffloadScheduler - from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( - KimiK3OffloadScheduler, - ) - from atom.kv_transfer.offload.metadata import LMCacheOffloadMetadata - monkeypatch.setattr( DenseOffloadScheduler, "build_connector_meta", lambda self: LMCacheOffloadMetadata(), ) - s = KimiK3OffloadScheduler.__new__(KimiK3OffloadScheduler) - s._pending_state_loads = [] - s._pending_state_stores = [] - s._pending_state_touches = [] - s._touch_round = 0 - s._state_touch_verdicts = {} - s._state_touches_unrefreshed = 0 + s = _k3_scheduler() s._refresh_save_stall = lambda: None return s @@ -7782,35 +7761,13 @@ def test_both_touch_reports_reach_quorum_across_tp(self, monkeypatch): ) merged = KVOutputAggregator(world_size=2).aggregate(outputs) s = self._scheduler(monkeypatch) + monkeypatch.setattr( + OffloadSchedulerMixin, "get_statistics", lambda self: {"load_requests": 3} + ) for completion in merged.connector_completions: assert s.connector_completion(completion) is None, "not a save" assert s.take_state_touch_verdicts() == {5: True} - assert s.take_state_touches_unrefreshed() == 1 - - def test_the_shells_forward_the_unrefreshed_count(self, monkeypatch): - from atom.kv_transfer.disaggregation.multi.multi_connector import ( - MultiConnectorScheduler, - ) - from atom.kv_transfer.offload.connector import ( - LMCacheOffloadConnectorScheduler, - ) - - shell = LMCacheOffloadConnectorScheduler.__new__( - LMCacheOffloadConnectorScheduler - ) - shell._impl = SimpleNamespace(take_state_touches_unrefreshed=lambda: 2) - assert shell.take_state_touches_unrefreshed() == 2 - shell._impl = SimpleNamespace() - assert shell.take_state_touches_unrefreshed() == 0 - - multi = MultiConnectorScheduler.__new__(MultiConnectorScheduler) - monkeypatch.setattr( - MultiConnectorScheduler, - "_state_tier_sub", - lambda self: SimpleNamespace(take_state_touches_unrefreshed=lambda: 4), - ) - assert multi.take_state_touches_unrefreshed() == 4 - monkeypatch.setattr( - MultiConnectorScheduler, "_state_tier_sub", lambda self: None - ) - assert multi.take_state_touches_unrefreshed() == 0 + assert s.get_statistics() == { + "load_requests": 3, + "state_touches_unrefreshed": 1, + } diff --git a/tests/test_state_checkpoint.py b/tests/test_state_checkpoint.py index b135ddbb69..84794c1def 100644 --- a/tests/test_state_checkpoint.py +++ b/tests/test_state_checkpoint.py @@ -17,6 +17,7 @@ from types import SimpleNamespace import pytest +from conftest import OFFLOAD_STALE_S as STALE from conftest import MockConfig from atom.model_engine.block_manager import ( @@ -34,7 +35,7 @@ from atom.model_engine.scheduler import ScheduledBatchOutput, Scheduler from atom.model_engine.sequence import Sequence, SequenceType from atom.model_engine.state_cache import StateCache -from atom.model_engine.state_offload import _TIER_VERIFIED_MAX_S, StateOffloadIndex +from atom.model_engine.state_offload import StateOffloadIndex from atom.model_engine.state_pool import StateSlotPool from atom.model_engine.state_runtime import ( StateRuntime, @@ -4566,24 +4567,12 @@ def test_a_cold_start_without_a_boundary_is_untouched(self): # ── state tier hit refresh: the admission hold and the voters ────────────── -@pytest.fixture -def clock(monkeypatch): - """`state_offload`'s monotonic clock, advanced by hand.""" - now = [1000.0] - monkeypatch.setattr("atom.model_engine.state_offload.monotonic", lambda: now[0]) - return now - - -#: Past every confirmation window `tier_verified_for_s` can give. -STALE = _TIER_VERIFIED_MAX_S + 1 - - def _paged_bm(): bm = make_block_manager( ckpt_config(num_kvcache_blocks=160, pool_entries={"state": 8}), state_runtime=PAGED_COPY_RUNTIME, ) - bm.state_offload = StateOffloadIndex(refresh_hits=True) + bm.state_offload = StateOffloadIndex(verified_for_s=5.0) return bm @@ -4611,11 +4600,11 @@ def _joint(self, bm, monkeypatch, *, h=77, from_hbm=False): ) return commits - def test_an_unconfirmed_joint_image_is_held(self, clock, monkeypatch): + def test_an_unconfirmed_joint_image_is_held(self, offload_clock, monkeypatch): bm = _paged_bm() commits = self._joint(bm, monkeypatch) bm.state_offload.note_stored(77) - clock[0] += STALE + offload_clock[0] += STALE seq = _seq() assert not bm.record_allocation( @@ -4630,13 +4619,13 @@ def test_an_unconfirmed_joint_image_is_held(self, clock, monkeypatch): ) assert len(commits) == 1 - def test_a_gone_image_is_not_parked_on(self, clock, monkeypatch): + def test_a_gone_image_is_not_parked_on(self, offload_clock, monkeypatch): """The verdict drops the hash; the retry's boundary search no longer sees the tier rung (`could_serve` is False), so it plans a lower one.""" bm = _paged_bm() self._joint(bm, monkeypatch) bm.state_offload.note_stored(77) - clock[0] += STALE + offload_clock[0] += STALE assert not bm.record_allocation( _seq(), AdmissionProbe(2, [1, 2]), defer_unverified=True ) @@ -4644,7 +4633,9 @@ def test_a_gone_image_is_not_parked_on(self, clock, monkeypatch): bm.state_offload.apply_touch_verdict(77, False) assert not bm._tier_can_serve(77) - def test_a_recently_confirmed_image_goes_straight_through(self, clock, monkeypatch): + def test_a_recently_confirmed_image_goes_straight_through( + self, offload_clock, monkeypatch + ): bm = _paged_bm() commits = self._joint(bm, monkeypatch) bm.state_offload.note_stored(77) @@ -4653,20 +4644,20 @@ def test_a_recently_confirmed_image_goes_straight_through(self, clock, monkeypat ) assert len(commits) == 1 - def test_a_state_resident_in_hbm_is_never_held(self, clock, monkeypatch): + def test_a_state_resident_in_hbm_is_never_held(self, offload_clock, monkeypatch): bm = _paged_bm() self._joint(bm, monkeypatch, from_hbm=True) bm.state_offload.note_stored(77) - clock[0] += STALE + offload_clock[0] += STALE assert bm.record_allocation( _seq(), AdmissionProbe(2, [1, 2]), defer_unverified=True ) - def test_the_fit_probe_path_never_holds(self, clock, monkeypatch): + def test_the_fit_probe_path_never_holds(self, offload_clock, monkeypatch): bm = _paged_bm() self._joint(bm, monkeypatch) bm.state_offload.note_stored(77) - clock[0] += STALE + offload_clock[0] += STALE assert bm.record_allocation(_seq(), AdmissionProbe(2, [1, 2])) def _hbm_hit(self, bm, monkeypatch, resident): @@ -4683,11 +4674,13 @@ def _hbm_hit(self, bm, monkeypatch, resident): ) bm.state_offload.note_stored(2) - def test_an_hbm_hit_holds_only_for_a_tier_sourced_state(self, clock, monkeypatch): + def test_an_hbm_hit_holds_only_for_a_tier_sourced_state( + self, offload_clock, monkeypatch + ): bm = _paged_bm() resident = {2} self._hbm_hit(bm, monkeypatch, resident) - clock[0] += STALE + offload_clock[0] += STALE assert bm.record_allocation( _seq(), AdmissionProbe(2, [1, 2]), defer_unverified=True ) @@ -4696,18 +4689,18 @@ def test_an_hbm_hit_holds_only_for_a_tier_sourced_state(self, clock, monkeypatch _seq(), AdmissionProbe(2, [1, 2]), defer_unverified=True ) - def test_a_replayed_hit_is_not_held(self, clock, monkeypatch): + def test_a_replayed_hit_is_not_held(self, offload_clock, monkeypatch): bm = _paged_bm() self._hbm_hit(bm, monkeypatch, set()) monkeypatch.setattr(bm, "state_replay_tokens", 8) - clock[0] += STALE + offload_clock[0] += STALE seq = _seq() seq.replays_hit = True assert bm.record_allocation( seq, AdmissionProbe(2, [1, 2]), defer_unverified=True ) - def test_the_hold_judges_the_joint_fit_kept(self, clock, monkeypatch): + def test_the_hold_judges_the_joint_fit_kept(self, offload_clock, monkeypatch): """A joint `_fit_joint` drops is not loaded, so its stale image must not hold an admission that resumes at an HBM-resident hit.""" bm = _paged_bm() @@ -4717,7 +4710,7 @@ def test_the_hold_judges_the_joint_fit_kept(self, clock, monkeypatch): ) monkeypatch.setattr(bm.paged_state_checkpoints, "contains", lambda h: h == 2) bm.state_offload.note_stored(77) - clock[0] += STALE + offload_clock[0] += STALE assert bm.record_allocation( _seq(), AdmissionProbe(2, [1, 2], 1), defer_unverified=True ) @@ -4734,12 +4727,12 @@ def _bm(self): bm.state_offload.note_stored(77) return bm - def test_the_resumable_scan_registers_a_tier_hit(self, clock): + def test_the_resumable_scan_registers_a_tier_hit(self, offload_clock): bm = self._bm() assert bm.paged_state_checkpoints._reachable(77) assert bm.state_offload.take_touches() == [77] - def test_an_hbm_hit_is_not_a_tier_vote(self, clock, monkeypatch): + def test_an_hbm_hit_is_not_a_tier_vote(self, offload_clock, monkeypatch): bm = self._bm() monkeypatch.setattr( bm.paged_state_checkpoints.store, "contains", lambda h: True @@ -4747,17 +4740,11 @@ def test_an_hbm_hit_is_not_a_tier_vote(self, clock, monkeypatch): assert bm.paged_state_checkpoints._reachable(77) assert bm.state_offload.take_touches() == [] - def test_asking_whether_the_tier_can_serve_registers_nothing(self, clock): + def test_asking_whether_the_tier_can_serve_registers_nothing(self, offload_clock): bm = self._bm() assert bm._tier_can_serve(77) assert bm.state_offload.take_touches() == [] - def test_an_offered_load_registers_its_hash(self, clock): - bm = self._bm() - seq = _seq() - assert bm._request_state_load(seq, 77) - assert bm.state_offload.take_touches() == [77] - class TestResumeHash: """`allocate` and the admission hold pick the state to resume from with @@ -4774,7 +4761,9 @@ def test_the_joint_boundary_wins_then_a_replay_then_the_hit(self, monkeypatch): assert bm._state_resume_hash(seq, -1, 2) == -1, "a replay restores none" assert bm._state_resume_hash(seq, 77, 2) == 77 - def test_the_hold_judges_the_hash_allocate_then_loads(self, clock, monkeypatch): + def test_the_hold_judges_the_hash_allocate_then_loads( + self, offload_clock, monkeypatch + ): bm = _paged_bm() offload = bm.state_offload bm.paged_state_checkpoints.attach_offload(offload) @@ -4787,12 +4776,12 @@ def test_the_hold_judges_the_hash_allocate_then_loads(self, clock, monkeypatch): # The image left HBM; only the tier has it, and not confirmed lately. bm.paged_state_checkpoints.unindex(h) offload.note_stored(h) - clock[0] += STALE + offload_clock[0] += STALE asked, attached = [], [] - verify, attach = offload.needs_verify, bm._attach_state_slots + hold, attach = offload.hold_admission, bm._attach_state_slots monkeypatch.setattr( - offload, "needs_verify", lambda x: asked.append(x) or verify(x) + offload, "hold_admission", lambda x: asked.append(x) or hold(x) ) monkeypatch.setattr( bm, "_attach_state_slots", lambda s, x: attached.append(x) or attach(s, x) @@ -4832,8 +4821,8 @@ def record(seq, probe, *, defer_unverified=False): assert seq.id in scheduled and not sched.waiting -def test_the_scheduler_publishes_touches_with_the_state_loads(clock): - idx = StateOffloadIndex(refresh_hits=True) +def test_the_scheduler_publishes_touches_with_the_state_loads(offload_clock): + idx = StateOffloadIndex(verified_for_s=5.0) idx.note_stored(7) idx.want_touch(7) sent = [] diff --git a/tests/test_state_offload_index.py b/tests/test_state_offload_index.py index 809d36c4bd..8c5cecc224 100644 --- a/tests/test_state_offload_index.py +++ b/tests/test_state_offload_index.py @@ -6,6 +6,7 @@ from types import SimpleNamespace import pytest +from conftest import OFFLOAD_STALE_S as STALE from atom.model_engine import state_offload as so from atom.model_engine.state_offload import ( @@ -62,51 +63,42 @@ def test_state_tier_cpu_bytes_reads_gib(monkeypatch): assert state_tier_cpu_bytes() == 96 << 30 -@pytest.fixture -def clock(monkeypatch): - """`state_offload`'s monotonic clock, advanced by hand.""" - now = [1000.0] - monkeypatch.setattr(so, "monotonic", lambda: now[0]) - return now - - # ── the index: votes queue touches, verdicts keep it in step ─────────────── -#: Past every confirmation window `tier_verified_for_s` can give. -STALE = so._TIER_VERIFIED_MAX_S + 1 - - -def _index(*hashes, refresh=True, **kw): - idx = StateOffloadIndex(refresh_hits=refresh, **kw) +def _index(*hashes, verified_for_s=so._TIER_VERIFIED_MIN_S, **kw): + idx = StateOffloadIndex(verified_for_s=verified_for_s, **kw) for h in hashes: idx.note_stored(h) return idx class TestIndexTouches: - def test_only_a_registered_hit_queues_a_touch(self, clock): + def test_only_a_registered_hit_queues_a_touch(self, offload_clock): idx = _index(7) assert idx.could_serve(7) assert idx.take_touches() == [], "could_serve is pure" + assert idx.request_load("r", 7) + assert idx.take_touches() == [], "a load is its own verdict" idx.want_touch(7) idx.want_touch(8) # not indexed: nothing to touch assert idx.take_touches() == [7] assert idx.take_touches() == [] - def test_no_touches_without_refresh(self, clock): - idx = _index(7, refresh=False) + def test_no_touches_without_refresh(self, offload_clock): + idx = _index(7, verified_for_s=None) idx.want_touch(7) assert idx.take_touches() == [] + assert idx._verified_at == {}, "nothing reads it" - def test_a_hash_voted_every_pass_is_touched_at_an_interval(self, clock): + def test_a_hash_voted_every_pass_is_touched_at_an_interval(self, offload_clock): idx = _index(7) idx.want_touch(7) assert idx.take_touches() == [7] - clock[0] += so._TIER_TOUCH_INTERVAL_S / 2 + offload_clock[0] += so._TIER_TOUCH_INTERVAL_S / 2 idx.want_touch(7) assert idx.take_touches() == [] - clock[0] += so._TIER_TOUCH_INTERVAL_S + offload_clock[0] += so._TIER_TOUCH_INTERVAL_S idx.want_touch(7) assert idx.take_touches() == [7] assert idx.touches_sent == 2 @@ -122,130 +114,144 @@ def test_a_hash_voted_every_pass_is_touched_at_an_interval(self, clock): ids=["gone", "stored-since", "restored-since", "stored-before"], ) def test_a_gone_verdict_drops_the_hash_unless_a_store_landed_since( - self, clock, before_touch, after_touch, dropped + self, offload_clock, before_touch, after_touch, dropped ): idx = _index(7) for op in before_touch: - clock[0] += 1 + offload_clock[0] += 1 getattr(idx, op)(7) - clock[0] += 1 + offload_clock[0] += 1 idx.want_touch(7) assert idx.take_touches() == [7] for op in after_touch: - clock[0] += 0.01 + offload_clock[0] += 0.01 getattr(idx, op)(7) idx.apply_touch_verdict(7, False) assert idx.could_serve(7) is not dropped assert idx.stats()["touch_misses"] == int(dropped) - def test_a_resident_verdict_refreshes_the_cap_order(self, clock): - idx = StateOffloadIndex(max_cpu_entries=19, refresh_hits=True) # cap 3 - for h in (1, 2, 3): - idx.note_stored(h) + def test_a_resident_verdict_refreshes_the_cap_order(self, offload_clock): + idx = _index(1, 2, 3, max_cpu_entries=19) # cap 3 idx.apply_touch_verdict(1, True) idx.note_stored(4) assert idx.could_serve(1) and not idx.could_serve(2) - def test_a_verdict_for_a_forgotten_hash_is_ignored(self, clock): + def test_a_verdict_for_a_forgotten_hash_is_ignored(self, offload_clock): idx = _index(7) idx.forget(7) idx.apply_touch_verdict(7, True) assert not idx.could_serve(7) - -class TestNeedsVerify: - def test_a_fresh_store_needs_no_verdict(self, clock): - assert not _index(7).needs_verify(7) - - def test_a_store_only_role_holds_nothing(self, clock): + def test_touch_state_leaves_with_its_hash(self, offload_clock): + idx = _index(1, 2, 3, max_cpu_entries=19) # cap 3 + for h in (1, 2, 3): + idx.want_touch(h) + idx.take_touches() + for h in (1, 2, 3): + idx.want_touch(h) + offload_clock[0] += STALE + assert idx.hold_admission(3) + idx.note_stored(4) + idx.note_stored(5) # evicts 1 and 2 + idx.forget(3) + for table in ( + idx._touch_wanted, + idx._touch_sent_at, + idx._verified_at, + idx._hold_started_at, + ): + assert set(table) <= idx.hashes + + +class TestHoldAdmission: + def test_a_fresh_store_needs_no_verdict(self, offload_clock): + assert not _index(7).hold_admission(7) + + def test_a_store_only_role_holds_nothing(self, offload_clock): """A `kv_producer` indexes its own stores but never loads them.""" idx = _index(7, can_load=False) - clock[0] += STALE - assert not idx.needs_verify(7) + offload_clock[0] += STALE + assert not idx.hold_admission(7) assert idx.take_touches() == [] - def test_a_stale_hash_waits_for_its_touch(self, clock): + def test_a_stale_hash_waits_for_its_touch(self, offload_clock): idx = _index(7) - clock[0] += STALE - assert idx.needs_verify(7) + offload_clock[0] += STALE + assert not idx.hold_admission(-1), "nothing to load" + assert idx.hold_admission(7) assert idx.take_touches() == [7], "the hold asks for its own touch" - assert idx.needs_verify(7), "still out" + assert idx.hold_admission(7), "still out" + assert idx.stats()["loads_deferred_for_verify"] == 2 idx.apply_touch_verdict(7, True) - assert not idx.needs_verify(7) + assert not idx.hold_admission(7) - def test_a_gone_verdict_ends_the_hold_by_dropping_the_hash(self, clock): + def test_a_gone_verdict_ends_the_hold_by_dropping_the_hash(self, offload_clock): idx = _index(7) - clock[0] += STALE - assert idx.needs_verify(7) + offload_clock[0] += STALE + assert idx.hold_admission(7) idx.take_touches() idx.apply_touch_verdict(7, False) - assert not idx.needs_verify(7) and not idx.could_serve(7) + assert not idx.hold_admission(7) and not idx.could_serve(7) - def test_a_lost_verdict_holds_no_longer_than_the_lookup_defer(self, clock): + def test_a_lost_verdict_holds_no_longer_than_the_lookup_defer(self, offload_clock): idx = _index(7) - clock[0] += STALE + offload_clock[0] += STALE # Re-probed every pass, re-touched every interval, never answered: # the re-sends must not extend the hold. held_for, step = 0.0, 0.05 - while idx.needs_verify(7): + while idx.hold_admission(7): assert held_for <= idx.defer_s + 1e-6 idx.take_touches() - clock[0] += step + offload_clock[0] += step held_for += step assert held_for >= idx.defer_s assert idx.touches_sent > 1, "the touch was re-sent while held" - assert not idx.needs_verify(7), "later admissions are not held again" + assert not idx.hold_admission(7), "later admissions are not held again" assert idx.stats()["verify_holds_expired"] == 1 - def test_evidence_after_an_expired_hold_rearms_it(self, clock): + def test_evidence_after_an_expired_hold_rearms_it(self, offload_clock): idx = _index(7) - clock[0] += STALE - assert idx.needs_verify(7) - clock[0] += idx.defer_s + 1 - assert not idx.needs_verify(7) + offload_clock[0] += STALE + assert idx.hold_admission(7) + offload_clock[0] += idx.defer_s + 1 + assert not idx.hold_admission(7) idx.apply_touch_verdict(7, True) - clock[0] += STALE - assert idx.needs_verify(7), "a fresh hold, timed from now" - clock[0] += idx.defer_s / 2 - assert idx.needs_verify(7) + offload_clock[0] += STALE + assert idx.hold_admission(7), "a fresh hold, timed from now" + offload_clock[0] += idx.defer_s / 2 + assert idx.hold_admission(7) - def test_a_verdict_ends_the_hold_timer(self, clock): + def test_a_verdict_ends_the_hold_timer(self, offload_clock): idx = _index(7) - clock[0] += STALE - assert idx.needs_verify(7) - clock[0] += idx.defer_s * 0.75 + offload_clock[0] += STALE + assert idx.hold_admission(7) + offload_clock[0] += idx.defer_s * 0.75 idx.apply_touch_verdict(7, True) - clock[0] += STALE - assert idx.needs_verify(7) - clock[0] += idx.defer_s * 0.75 - assert idx.needs_verify(7), "the earlier wait does not count" + offload_clock[0] += STALE + assert idx.hold_admission(7) + offload_clock[0] += idx.defer_s * 0.75 + assert idx.hold_admission(7), "the earlier wait does not count" - def test_the_lookup_defer_is_read_once_at_construction(self, clock, monkeypatch): + def test_the_lookup_defer_is_read_once_at_construction( + self, offload_clock, monkeypatch + ): monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "0") idx = _index(7) - clock[0] += STALE - assert not idx.needs_verify(7), "0 disables the hold" + offload_clock[0] += STALE + assert not idx.hold_admission(7), "0 disables the hold" monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "2") - assert not idx.needs_verify(7), "not re-read on the admission path" + assert not idx.hold_admission(7), "not re-read on the admission path" def test_an_index_without_refresh_never_reads_the_defer(self, monkeypatch): monkeypatch.setenv("OFFLOAD_LOOKUP_DEFER_S", "not-a-number") assert StateOffloadIndex().defer_s == 0 - def test_a_completed_load_counts_as_verified(self, clock): + def test_a_completed_load_counts_as_verified(self, offload_clock): idx = _index(7) - clock[0] += STALE + offload_clock[0] += STALE assert idx.request_load("r", 7) idx.complete_load("r") - assert not idx.needs_verify(7) - - def test_a_hold_is_counted(self, clock): - idx = _index(7) - assert not idx.hold_admission(-1), "nothing to load" - assert not idx.hold_admission(7), "confirmed by its store" - clock[0] += STALE - assert idx.hold_admission(7) - assert idx.stats()["loads_deferred_for_verify"] == 1 + assert not idx.hold_admission(7) @pytest.mark.parametrize( @@ -255,19 +261,19 @@ def test_the_verified_window_scales_with_the_pool_within_bounds(pool_gib, window assert so.tier_verified_for_s(pool_gib) == pytest.approx(window) -def test_the_index_holds_by_its_own_window(clock): +def test_the_index_holds_by_its_own_window(offload_clock): idx = _index(7, verified_for_s=8.0) - clock[0] += 7.5 - assert not idx.needs_verify(7) - clock[0] += 1 - assert idx.needs_verify(7) + offload_clock[0] += 7.5 + assert not idx.hold_admission(7) + offload_clock[0] += 1 + assert idx.hold_admission(7) class TestTouchPlumbing: """The scheduler's two calls a pass: `publish_touches` before the metadata is built, `collect_touch_verdicts` before the next admission.""" - def test_voted_hashes_reach_the_connector(self, clock): + def test_voted_hashes_reach_the_connector(self, offload_clock): idx = _index(7) idx.want_touch(7) sent = [] @@ -276,44 +282,37 @@ def test_voted_hashes_reach_the_connector(self, clock): ) assert sent == [7] and idx.refresh_hits - def test_a_connector_without_touches_turns_the_refresh_off(self, clock): + def test_a_connector_without_touches_turns_the_refresh_off(self, offload_clock): idx = _index(7) idx.want_touch(7) idx.publish_touches(SimpleNamespace()) idx.want_touch(7) assert idx.take_touches() == [] - clock[0] += STALE - assert not idx.needs_verify(7), "nothing would ever answer the hold" + offload_clock[0] += STALE + assert not idx.hold_admission(7), "nothing would ever answer the hold" - def test_verdicts_and_presence_only_touches_are_collected(self, clock): + def test_verdicts_are_collected(self, offload_clock): idx = _index(7, 8) - clock[0] += 1 + offload_clock[0] += 1 idx.want_touch(7) idx.want_touch(8) assert sorted(idx.take_touches()) == [7, 8] - connector = SimpleNamespace( - take_state_touch_verdicts=lambda: {7: False, 8: True}, - take_state_touches_unrefreshed=lambda: 3, + idx.collect_touch_verdicts( + SimpleNamespace(take_state_touch_verdicts=lambda: {7: False, 8: True}) ) - idx.collect_touch_verdicts(connector) - idx.collect_touch_verdicts(connector) assert not idx.could_serve(7) and idx.could_serve(8) - assert idx.stats()["touches_unrefreshed"] == 6 - - -GIB = 1 << 30 @pytest.mark.parametrize( - "state, pool, says", + "state_gib, pool_gib, says", [ (59, 208, None), (32, 32, "bounds nothing"), # #2562's accuracy run: the default cap (40, 64, "more than half"), + (32, 64, None), (32, 0, None), - (0, 32, None), ], ) -def test_the_state_cap_is_judged_against_the_pool(state, pool, says): - msg = state_tier_budget_warning(state * GIB, pool * GIB) +def test_the_state_cap_is_judged_against_the_pool(state_gib, pool_gib, says): + msg = state_tier_budget_warning(state_gib, pool_gib) assert msg is None if says is None else says in msg From caf4ea8f5c1aabaf1ef353cf96591999546d22f3 Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Mon, 2 Nov 2026 05:09:39 +0000 Subject: [PATCH 14/18] tests: keep the state-touch tests off the file's tail Appended after the last test they collided with #2636's appended tests on a merge; the block moves above TestStateLoadsAndStoresRunInSeparateLanes. --- tests/test_lmcache_offload_connector.py | 332 ++++++++++++------------ 1 file changed, 166 insertions(+), 166 deletions(-) diff --git a/tests/test_lmcache_offload_connector.py b/tests/test_lmcache_offload_connector.py index 4cfddc9011..59c181f5a4 100644 --- a/tests/test_lmcache_offload_connector.py +++ b/tests/test_lmcache_offload_connector.py @@ -7188,172 +7188,6 @@ def test_the_scheduler_half_keeps_the_operation_whole(self): # ── kimi_k3: loads must not queue behind a backlog of stores ─────────────── -class TestStateLoadsAndStoresRunInSeparateLanes: - """A load is on the TTFT critical path and a store is not, but one serial - executor made that ordering unenforceable: a load submitted in a later step - sat behind every store already queued, and a stuck store blocked all of - them. Putting same-step loads first cannot overtake queued work.""" - - class _BlockingCodec: - """A codec whose stores hang until released; loads always land.""" - - def __init__(self): - self.gate = threading.Event() - self.loaded = threading.Event() - - def put(self, h, unit_ids, on_source_released=None): - self.gate.wait(timeout=5) - if on_source_released is not None: - on_source_released() - return True - - def get(self, h, slot): - self.loaded.set() - return True - - def _tier(self, codec, **kw): - from atom.kv_transfer.offload.hybrid.kimi_k3.state_tier import ( - StateOffloadTier, - ) - - return StateOffloadTier(codec, **kw) - - def test_a_load_lands_while_a_store_is_still_stuck(self): - codec = self._BlockingCodec() - tier = self._tier(codec) - try: - for gen in range(4): - tier.submit_store(StateStoreOperationId(gen, gen + 1), (0,)) - tier.submit_load("r1", 77, 0) - assert codec.loaded.wait( - timeout=5 - ), "the load waited behind the store backlog" - done, failed = tier.get_finished() - assert done == {"r1"} - assert failed == set() - finally: - codec.gate.set() - tier.shutdown() - - def test_one_lane_serialises_the_load_behind_the_inflight_store(self): - """`staging_lanes=1` serialises the two lanes: a load waits out the - single in-flight store, but still not the backlog behind it. It does - *not* change standing HBM -- the staging buffer is per-thread and both - executors are `max_workers=1`, so HBM is two buffers either way; this - knob only gates load/store concurrency, which is what this asserts.""" - codec = self._BlockingCodec() - tier = self._tier(codec, staging_lanes=1) - try: - for gen in range(3): - tier.submit_store(StateStoreOperationId(gen, gen + 1), (0,)) - tier.submit_load("r1", 77, 0) - assert not codec.loaded.wait(timeout=0.2), "held by the in-flight store" - codec.gate.set() - assert codec.loaded.wait(timeout=5) - finally: - codec.gate.set() - tier.shutdown() - - def test_the_oldest_store_age_is_visible_while_it_hangs(self): - codec = self._BlockingCodec() - tier = self._tier(codec) - try: - assert tier.oldest_store_age_s() == 0.0 - tier.submit_store(StateStoreOperationId(1, 1), (0,)) - assert tier.oldest_store_age_s() >= 0.0 - assert len(tier._store_submitted_at) == 1 - finally: - codec.gate.set() - tier.drain() - assert tier.oldest_store_age_s() == 0.0 - tier.shutdown() - - class _EarlyReleaseCodec: - """`put` releases the source (D2H drained) and then stalls in the CPU - `batched_put`, so a drain can observe the early release *before* the - store completes -- exactly the window the source/store split exists - for.""" - - def __init__(self): - self.released = threading.Event() - self.gate = threading.Event() - - def put(self, h, unit_ids, on_source_released=None): - if on_source_released is not None: - on_source_released() # D2H is done here, CPU put is not - self.released.set() - self.gate.wait(timeout=5) - return True - - def get(self, h, slot): - return True - - def test_source_release_is_emitted_once_early_not_again_at_store_end(self): - """The release must land the instant the D2H drains (before the CPU - put), and must not be re-emitted when `_do_store` returns. A second - emission across two engine drains double-unpins the same record.""" - codec = self._EarlyReleaseCodec() - tier = self._tier(codec) - op = StateStoreOperationId(1, 1) - try: - tier.submit_store(op, (0,)) - assert codec.released.wait(timeout=5), "the D2H callback never fired" - # Early: units handed back while the CPU put is still running. - assert tier.take_source_releases() == {op} - # Finish the store; the end-of-store backstop must stay silent. - codec.gate.set() - tier.drain() - assert tier.take_source_releases() == set(), "source release re-emitted" - assert tier.take_store_reports() == ({op}, set()) - finally: - codec.gate.set() - tier.shutdown() - - def test_a_refused_store_still_releases_its_units(self): - """A refused allocation never fires the D2H callback and never read the - units, so the backstop must release them -- withholding it would hold an - image out of the KV pool until the stale reclaimer noticed.""" - - class _RefuseCodec: - def put(self, h, unit_ids, on_source_released=None): - return False # allocation refused; callback never reached - - def get(self, h, slot): - return True - - tier = self._tier(_RefuseCodec()) - op = StateStoreOperationId(2, 1) - try: - tier.submit_store(op, (0,)) - tier.drain() - assert tier.take_source_releases() == {op} - assert tier.take_store_reports() == (set(), {op}) - finally: - tier.shutdown() - - def test_a_throwing_store_still_releases_its_units(self): - """A throwing `pack` drains the device before it propagates, so the GPU - has stopped reading the units by the time control returns -- the backstop - releases them and the store is reported failed, not lost.""" - - class _ThrowCodec: - def put(self, h, unit_ids, on_source_released=None): - raise RuntimeError("pack blew up") - - def get(self, h, slot): - return True - - tier = self._tier(_ThrowCodec()) - op = StateStoreOperationId(3, 1) - try: - tier.submit_store(op, (0,)) - tier.drain() - assert tier.take_source_releases() == {op} - assert tier.take_store_reports() == (set(), {op}) - finally: - tier.shutdown() - - class _Obj: """A MemoryObj's reference count, and nothing else.""" @@ -7771,3 +7605,169 @@ def test_both_touch_reports_reach_quorum_across_tp(self, monkeypatch): "load_requests": 3, "state_touches_unrefreshed": 1, } + + +class TestStateLoadsAndStoresRunInSeparateLanes: + """A load is on the TTFT critical path and a store is not, but one serial + executor made that ordering unenforceable: a load submitted in a later step + sat behind every store already queued, and a stuck store blocked all of + them. Putting same-step loads first cannot overtake queued work.""" + + class _BlockingCodec: + """A codec whose stores hang until released; loads always land.""" + + def __init__(self): + self.gate = threading.Event() + self.loaded = threading.Event() + + def put(self, h, unit_ids, on_source_released=None): + self.gate.wait(timeout=5) + if on_source_released is not None: + on_source_released() + return True + + def get(self, h, slot): + self.loaded.set() + return True + + def _tier(self, codec, **kw): + from atom.kv_transfer.offload.hybrid.kimi_k3.state_tier import ( + StateOffloadTier, + ) + + return StateOffloadTier(codec, **kw) + + def test_a_load_lands_while_a_store_is_still_stuck(self): + codec = self._BlockingCodec() + tier = self._tier(codec) + try: + for gen in range(4): + tier.submit_store(StateStoreOperationId(gen, gen + 1), (0,)) + tier.submit_load("r1", 77, 0) + assert codec.loaded.wait( + timeout=5 + ), "the load waited behind the store backlog" + done, failed = tier.get_finished() + assert done == {"r1"} + assert failed == set() + finally: + codec.gate.set() + tier.shutdown() + + def test_one_lane_serialises_the_load_behind_the_inflight_store(self): + """`staging_lanes=1` serialises the two lanes: a load waits out the + single in-flight store, but still not the backlog behind it. It does + *not* change standing HBM -- the staging buffer is per-thread and both + executors are `max_workers=1`, so HBM is two buffers either way; this + knob only gates load/store concurrency, which is what this asserts.""" + codec = self._BlockingCodec() + tier = self._tier(codec, staging_lanes=1) + try: + for gen in range(3): + tier.submit_store(StateStoreOperationId(gen, gen + 1), (0,)) + tier.submit_load("r1", 77, 0) + assert not codec.loaded.wait(timeout=0.2), "held by the in-flight store" + codec.gate.set() + assert codec.loaded.wait(timeout=5) + finally: + codec.gate.set() + tier.shutdown() + + def test_the_oldest_store_age_is_visible_while_it_hangs(self): + codec = self._BlockingCodec() + tier = self._tier(codec) + try: + assert tier.oldest_store_age_s() == 0.0 + tier.submit_store(StateStoreOperationId(1, 1), (0,)) + assert tier.oldest_store_age_s() >= 0.0 + assert len(tier._store_submitted_at) == 1 + finally: + codec.gate.set() + tier.drain() + assert tier.oldest_store_age_s() == 0.0 + tier.shutdown() + + class _EarlyReleaseCodec: + """`put` releases the source (D2H drained) and then stalls in the CPU + `batched_put`, so a drain can observe the early release *before* the + store completes -- exactly the window the source/store split exists + for.""" + + def __init__(self): + self.released = threading.Event() + self.gate = threading.Event() + + def put(self, h, unit_ids, on_source_released=None): + if on_source_released is not None: + on_source_released() # D2H is done here, CPU put is not + self.released.set() + self.gate.wait(timeout=5) + return True + + def get(self, h, slot): + return True + + def test_source_release_is_emitted_once_early_not_again_at_store_end(self): + """The release must land the instant the D2H drains (before the CPU + put), and must not be re-emitted when `_do_store` returns. A second + emission across two engine drains double-unpins the same record.""" + codec = self._EarlyReleaseCodec() + tier = self._tier(codec) + op = StateStoreOperationId(1, 1) + try: + tier.submit_store(op, (0,)) + assert codec.released.wait(timeout=5), "the D2H callback never fired" + # Early: units handed back while the CPU put is still running. + assert tier.take_source_releases() == {op} + # Finish the store; the end-of-store backstop must stay silent. + codec.gate.set() + tier.drain() + assert tier.take_source_releases() == set(), "source release re-emitted" + assert tier.take_store_reports() == ({op}, set()) + finally: + codec.gate.set() + tier.shutdown() + + def test_a_refused_store_still_releases_its_units(self): + """A refused allocation never fires the D2H callback and never read the + units, so the backstop must release them -- withholding it would hold an + image out of the KV pool until the stale reclaimer noticed.""" + + class _RefuseCodec: + def put(self, h, unit_ids, on_source_released=None): + return False # allocation refused; callback never reached + + def get(self, h, slot): + return True + + tier = self._tier(_RefuseCodec()) + op = StateStoreOperationId(2, 1) + try: + tier.submit_store(op, (0,)) + tier.drain() + assert tier.take_source_releases() == {op} + assert tier.take_store_reports() == (set(), {op}) + finally: + tier.shutdown() + + def test_a_throwing_store_still_releases_its_units(self): + """A throwing `pack` drains the device before it propagates, so the GPU + has stopped reading the units by the time control returns -- the backstop + releases them and the store is reported failed, not lost.""" + + class _ThrowCodec: + def put(self, h, unit_ids, on_source_released=None): + raise RuntimeError("pack blew up") + + def get(self, h, slot): + return True + + tier = self._tier(_ThrowCodec()) + op = StateStoreOperationId(3, 1) + try: + tier.submit_store(op, (0,)) + tier.drain() + assert tier.take_source_releases() == {op} + assert tier.take_store_reports() == (set(), {op}) + finally: + tier.shutdown() From 0d5600b396b7864faa3dd6f4b6373e0f14d740f0 Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Mon, 2 Nov 2026 09:16:28 +0000 Subject: [PATCH 15/18] Keep #2668 to the touch and hold: revert edits to existing code Reverts the hunks that reworded, reflowed or refactored code the feature does not need to change; the behaviour is unchanged. - types.py: StateStoreOperationId docstring back to main. - _offload_common.py: no new StateOffloadFace abstract methods (the shell and multi forwards already use getattr). - kimi_k3/state_object.py: drop the _cache_keys comment and the get() miss-pop; get() keeps only the LRU refresh. put() keeps its original key line. - kimi_k3/connector.py: max_cache_bytes=state_tier_cpu_bytes() as on main. - state_offload.py: original class, forget, could_serve and state_tier_cpu_bytes docstrings; state_tier_chunk_tokens no longer refactored (state_tier_verified_for_s reads the pool size itself). - block_manager.py: record_allocation keeps its original shape, allocate is untouched; _resume_hash_off_hbm mirrors allocate's choice instead of sharing a new _state_resume_hash. - scheduler.py: original _publish_state_loads docstring and lines; the touch calls are inserted only. - tests: conftest.py and _FakeStorage back to main; the clock fixture lives in the two files that use it; original test module docstring restored; the new block sits above the existing section header. --- atom/kv_transfer/disaggregation/types.py | 9 +- atom/kv_transfer/offload/README.md | 5 +- atom/kv_transfer/offload/_offload_common.py | 6 +- .../offload/hybrid/kimi_k3/connector.py | 7 +- .../offload/hybrid/kimi_k3/state_object.py | 12 +- atom/model_engine/block_manager.py | 105 ++++++++---------- atom/model_engine/scheduler.py | 19 ++-- atom/model_engine/state_offload.py | 65 ++++++----- tests/conftest.py | 13 --- tests/test_block_manager.py | 5 +- tests/test_lmcache_offload_connector.py | 36 ++---- tests/test_state_checkpoint.py | 29 +++-- tests/test_state_offload_index.py | 17 ++- 13 files changed, 140 insertions(+), 188 deletions(-) diff --git a/atom/kv_transfer/disaggregation/types.py b/atom/kv_transfer/disaggregation/types.py index a842f3fb45..bfb6393a8f 100644 --- a/atom/kv_transfer/disaggregation/types.py +++ b/atom/kv_transfer/disaggregation/types.py @@ -109,11 +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 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. + 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. """ prefix_hash: int diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index 2c7492e7ea..d4187713d1 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -131,9 +131,8 @@ 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`. `OFFLOAD_STATE_CPU_SIZE` -caps the images' share of the pool and reserves nothing; the worker warns when it -is more than half of the pool. +the connector statistic `state_touches_unrefreshed`. The worker warns when +`OFFLOAD_STATE_CPU_SIZE` is more than half of the pool. ## LMCache Multiprocess (`lmcache_mp`) diff --git a/atom/kv_transfer/offload/_offload_common.py b/atom/kv_transfer/offload/_offload_common.py index a5e70ba0ca..cd26767a23 100644 --- a/atom/kv_transfer/offload/_offload_common.py +++ b/atom/kv_transfer/offload/_offload_common.py @@ -426,7 +426,7 @@ class StateOffloadFace(ABC): `QwenOffloadScheduler` inherits that, so `isinstance(impl, StateOffloadFace)` is the honest predicate. - Narrow on purpose -- only the tier methods. Other impl attributes the + Narrow on purpose -- only the four tier methods. Other impl attributes the scheduler reads (e.g. `max_pending_saves`) stay plain shell forwards rather than joining this face, which is exclusively the KDA state-tier contract. """ @@ -439,10 +439,6 @@ def enqueue_state_stores(self, stores) -> bool: ... def take_state_reports(self) -> tuple[set[int], set[int]]: ... @abstractmethod def take_state_source_releases(self) -> set: ... - @abstractmethod - def enqueue_state_touches(self, hashes) -> bool: ... - @abstractmethod - def take_state_touch_verdicts(self) -> dict[int, bool]: ... class OffloadSchedulerMixin(ABC): diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py index b839c5716a..190b446664 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py @@ -249,7 +249,6 @@ def _build_state_tier(self, transfer_tensors) -> None: staging_buffer_bytes=entry_bytes, release_after_transfer=gpu_connector.release_gpu_staging_after_transfer, ) - state_bytes = state_tier_cpu_bytes() codec = StateByteCodec( backend, staged, @@ -258,13 +257,15 @@ def _build_state_tier(self, transfer_tensors) -> None: world_size=world, worker_id=rank, layout_id=layout_id, - max_cache_bytes=state_bytes, + max_cache_bytes=state_tier_cpu_bytes(), ) # Share the allocator, but cap state retention separately so old # checkpoints cannot consume nearly the whole paged-KV CPU pool. codec.bind_storage_manager(self._engine.storage_manager) pool_gib = float(getattr(cfg, "max_local_cpu_size", 0) or 0) - budget_warning = state_tier_budget_warning(state_bytes / (1 << 30), pool_gib) + budget_warning = state_tier_budget_warning( + state_tier_cpu_bytes() / (1 << 30), pool_gib + ) if budget_warning: logger.warning("kimi_k3 offload: %s", budget_warning) # No index here: StateOffloadIndex lives in the engine process; both diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py index 83a2e8fbdb..4ec6da918a 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py @@ -106,7 +106,6 @@ def __init__( ) if self._max_entries is not None and self._max_entries < 1: raise ValueError("state CPU budget must fit at least one checkpoint") - # The cap's LRU, kept in LMCache's order (`get`, `touch`, a re-store). self._cache_keys: OrderedDict[Any, None] = OrderedDict() self._cache_lock = threading.Lock() self._store_lock = threading.Lock() @@ -164,8 +163,7 @@ def put(self, h: int, unit_ids, on_source_released=None) -> bool: # unpacking, so removing its cache entry cannot free live bytes. if self._storage is None: return False - key = self.key(h) - if self._touch_key(key) is TouchResult.REFRESHED: + if self._touch_key(self.key(h)) is TouchResult.REFRESHED: # Same key, same bytes, already hot: a re-store is a D2H LocalCPU # refuses as a duplicate. One held only by disk is stored again. if on_source_released is not None: @@ -173,6 +171,7 @@ def put(self, h: int, unit_ids, on_source_released=None) -> bool: return True if self._max_entries is None: return self._put(h, unit_ids, on_source_released) + key = self.key(h) with self._store_lock: with self._cache_lock: if key not in self._cache_keys: @@ -309,11 +308,8 @@ def get(self, h: int, slot) -> bool: """ if self._storage is None: return False - key = self.key(h) - obj = self._storage.get(key) + obj = self._storage.get(self.key(h)) if obj is None: - with self._cache_lock: - self._cache_keys.pop(key, None) return False size = self._object_bytes(obj) if size is not None and size != self.entry_bytes: @@ -338,7 +334,7 @@ def get(self, h: int, slot) -> bool: finally: obj.ref_count_down() # `get_blocking` does not refresh LMCache's LRU; a resume is a use. - self._touch_key(key) + self._touch_key(self.key(h)) return True @staticmethod diff --git a/atom/model_engine/block_manager.py b/atom/model_engine/block_manager.py index 4f8c2d168e..7ddc310cc7 100644 --- a/atom/model_engine/block_manager.py +++ b/atom/model_engine/block_manager.py @@ -1552,23 +1552,22 @@ def record_allocation( restore from outside HBM (`StateOffloadIndex.hold_admission`). The caller retries the request a pass later. """ - if not self.enable_prefix_caching: - return True - decision = self._joint_kv_boundary( - seq, probe.num_cached_blocks, probe.block_hashes - ) - decision = self._fit_joint(decision, probe) - if ( - defer_unverified - and self.state_offload is not None - and self.state_offload.hold_admission( - self._resume_hash_off_hbm( - seq, probe.num_cached_blocks, probe.block_hashes, decision - ) + if self.enable_prefix_caching: + decision = self._joint_kv_boundary( + seq, probe.num_cached_blocks, probe.block_hashes ) - ): - return False - self._commit_joint_boundary(seq, decision) + decision = self._fit_joint(decision, probe) + if ( + defer_unverified + and self.state_offload is not None + and self.state_offload.hold_admission( + self._resume_hash_off_hbm( + seq, probe.num_cached_blocks, probe.block_hashes, decision + ) + ) + ): + return False + self._commit_joint_boundary(seq, decision) return True def _resume_hash_off_hbm( @@ -1578,41 +1577,28 @@ def _resume_hash_off_hbm( block_hashes: list[int], decision: JointBoundaryDecision, ) -> int: - """The `_state_resume_hash` `allocate` would restore from, when HBM - does not hold it; -1 for no state, an HBM-held one, or a replay.""" + """The hash `allocate` would restore `seq`'s state from, when HBM does + not hold it; -1 for no state, an HBM-held one, or a replay. + + Mirrors `allocate`'s choice: the joint boundary when there is one (above + the HBM hit by construction), else the HBM hit's last block unless a + replay rebuilds the state instead. + """ if not seq.has_per_req_cache: return -1 - joint = decision.skip_reason is None - if joint and decision.state_from_hbm: + if decision.skip_reason is None: + return -1 if decision.state_from_hbm else decision.boundary_hash + if self.tier_resume_replays(seq) or not ( + 0 < num_cached_blocks <= len(block_hashes) + ): return -1 - hit_hash = ( - block_hashes[num_cached_blocks - 1] - if 0 < num_cached_blocks <= len(block_hashes) - else -1 - ) - h = self._state_resume_hash( - seq, decision.boundary_hash if joint else -1, hit_hash + h = block_hashes[num_cached_blocks - 1] + in_hbm = ( + self.paged_state_checkpoints.contains(h) + if self.paged_state_checkpoints is not None + else self.state.lookup(h) >= 0 ) - if h != -1 and not joint: - in_hbm = ( - self.paged_state_checkpoints.contains(h) - if self.paged_state_checkpoints is not None - else self.state.lookup(h) >= 0 - ) - if in_hbm: - return -1 - return h - - def _state_resume_hash(self, seq: Sequence, joint_hash: int, hit_hash: int) -> int: - """The hash whose state `seq` resumes from, or -1 for none. - - The joint boundary when there is one (above the HBM hit by - construction), else the HBM hit's last block. A replay rebuilds the - state instead of restoring it, so it resumes from none. - """ - if joint_hash != -1: - return joint_hash - return -1 if self.tier_resume_replays(seq) else hit_hash + return -1 if in_hbm else h def retreat_joint_to_hbm(self, seq: Sequence) -> bool: """Drop a failed unaligned joint and resume at the HBM hit it lowered. @@ -1809,16 +1795,13 @@ def allocate(self, seq: Sequence, num_cached_blocks: int = 0) -> bool: # last one actually claimed, whichever way the loop exits. if i < num_cached_blocks: hit_hash = h - # A replay starts from a fresh state slot, never a checkpoint: it - # rebuilds the state below whatever the claim above reached - # (`_state_resume_hash` gives it none). A KV-only load past the hit - # would start above a gap. - if ( - self.state_replay_tokens - and seq.replays_hit - and num_cached_blocks < requested_blocks - ): - seq.replay_kv_load = False + if self.state_replay_tokens and seq.replays_hit: + # A replay starts from a fresh state slot, never a checkpoint: it + # rebuilds the state below whatever the claim above reached. + hit_hash = -1 + if num_cached_blocks < requested_blocks: + # A KV-only load past the hit would start above a gap. + seq.replay_kv_load = False # Pin the restore before fresh blocks can evict its checkpoint. state_holds = True if seq.has_per_req_cache and self.paged_state_checkpoints is not None: @@ -1834,9 +1817,9 @@ def allocate(self, seq: Sequence, num_cached_blocks: int = 0) -> bool: # joint boundary and `_claim_after_load` raised `num_cached_tokens` # to it, so the forward resumed over a prefix the state does not # cover. Silent wrong output, no exception. + joint_hash = seq.offload_joint.boundary_hash state_holds = self._attach_state_slots( - seq, - self._state_resume_hash(seq, seq.offload_joint.boundary_hash, hit_hash), + seq, joint_hash if joint_hash != -1 else hit_hash ) for _ in range(len(seq.block_table), self.num_pool_blocks(len(seq))): seq.block_table.append(self._fresh_block()) @@ -1853,9 +1836,9 @@ def allocate(self, seq: Sequence, num_cached_blocks: int = 0) -> bool: # above the HBM hit by construction. `num_cached_tokens` stays at # the HBM prefix until the KV leg lands, so the forward covers # `[hbm, num_prompt)` if anything goes wrong from here. + joint_hash = seq.offload_joint.boundary_hash state_holds = self._attach_state_slots( - seq, - self._state_resume_hash(seq, seq.offload_joint.boundary_hash, hit_hash), + seq, joint_hash if joint_hash != -1 else hit_hash ) if seq.has_per_req_cache: seq._state_initialized_after_alloc = False diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index 30f62c6b5a..c25e633a64 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -3779,8 +3779,8 @@ def _publish_state_stores(self) -> None: self.block_manager.settle_state_store(op, ok=False, attempted=False) def _publish_state_loads(self) -> None: - """Hand this pass's state-tier loads, and the touches the admission - pass registered, to the connector before it builds its metadata. + """Hand this pass's state-tier loads to the connector, before it builds + its metadata. Drained here rather than at the park, because the connector publishes once per pass and a load handed over twice would write the same entry @@ -3791,11 +3791,10 @@ def _publish_state_loads(self) -> None: # Guarded like the `_settle_state_load` / `_publish_state_stores` # siblings: a scheduler built without a block manager (the connector # test doubles) would otherwise AttributeError on the bare deref. - bm = getattr(self, "block_manager", None) - offload = getattr(bm, "state_offload", None) + offload = getattr(getattr(self, "block_manager", None), "state_offload", None) if offload is not None: offload.publish_touches(self.kv_connector) - take = getattr(bm, "take_state_loads", None) + take = getattr(getattr(self, "block_manager", None), "take_state_loads", None) if take is None: return loads = take() @@ -4907,11 +4906,6 @@ def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput): # offload leg. self._drain_source_safe_releases() - bm = getattr(self, "block_manager", None) - offload = getattr(bm, "state_offload", None) - if offload is not None: - # Touch verdicts, before the next admission pass reads them. - offload.collect_touch_verdicts(self.kv_connector) # The state-offload store reports, not keyed by request: by the time a # store lands its owner is long gone and only the hash remains. They # ride the connector's completion channels, so they are drained from the @@ -4920,6 +4914,11 @@ def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput): # Indexed only on the report, never at submission: a hash advertised # before its bytes exist parks the next request over that prefix # against a `get` that must miss. + bm = getattr(self, "block_manager", None) + offload = getattr(bm, "state_offload", None) + if offload is not None: + # Touch verdicts, before the next admission pass reads them. + offload.collect_touch_verdicts(self.kv_connector) take = getattr(self.kv_connector, "take_state_reports", None) if offload is not None and take is not None: indexed, failed = take() diff --git a/atom/model_engine/state_offload.py b/atom/model_engine/state_offload.py index 39b98f8c92..2c50e62ae3 100644 --- a/atom/model_engine/state_offload.py +++ b/atom/model_engine/state_offload.py @@ -35,15 +35,15 @@ class StateOffloadIndex: """What is believed to be in LMCache, and what is being fetched back. - `hashes` answers the membership half of the tier's vote. LMCache's own LRU - can drop bytes at any time, so on its own a hash here means "was stored - once", never "is still there". + `hashes` answers the membership half of the tier's vote. It is deliberately + optimistic: LMCache's own LRU can drop bytes at any time, so a hash here + means "was stored once", never "is still there". The false positive costs + one lookup and a park/unpark, and is handled by the `failed_loading` path + (`fail_load` -> `forget`). With a `verified_for_s` window, every voted hit is also touched: the worker refreshes the image in LMCache's LRU, as a KV lookup does for its chunks, - and reports whether it was still there (`apply_touch_verdict`). A false - positive that slips through still costs one park/unpark and is handled by - the `failed_loading` path (`fail_load` -> `forget`). + and reports whether it was still there (`apply_touch_verdict`). """ def __init__( @@ -147,7 +147,7 @@ def note_stored(self, h: int) -> None: self.hashes_evicted += 1 def forget(self, h: int) -> None: - """Drop a hash whose image is gone, so the next request does not retry.""" + """Drop a hash whose load failed, so the next request does not retry.""" self.hashes.discard(h) self._hash_lru.pop(h, None) self._drop_touch_state(h) @@ -169,8 +169,7 @@ def _end_hold(self, h: int) -> None: def could_serve(self, h: int) -> bool: """Whether a load for `h` could be offered at all: this tier can load - AND believes it holds `h`. Pure: a voter registers its hit with - `want_touch`. + AND believes it holds `h`. The membership+capability half of `request_load`'s guard (minus its per-request in-flight check), factored out so the admission-path voters @@ -567,8 +566,7 @@ def state_tier_capability(config) -> StateTierCapability: def state_tier_cpu_bytes() -> int: """CPU bytes the state tier may keep per worker (`OFFLOAD_STATE_CPU_SIZE`, - GiB, default 32): a cap on the images in the paged-KV LMCache pool they - share, not a reservation -- images and KV chunks compete in one LRU. + GiB, default 32), carved from the paged-KV LMCache pool it shares. One reader for both halves: the worker's `StateByteCodec` evicts against it and the engine's `StateOffloadIndex` mirrors the resulting entry bound. @@ -600,27 +598,6 @@ def state_tier_budget_warning(state_gib: float, pool_gib: float) -> str | None: return None -def _lmcache_engine_config(config): - """The `LMCacheEngineConfig` this engine's offload connector builds. - - Raises when LMCache or its config cannot be read; the callers decide what - that costs. - """ - from atom.kv_transfer.offload.config import build_lmcache_config - - kvcfg = getattr(config, "kv_transfer_config", None) or {} - # Under `multi` the lmcache.* keys (chunk_size, offload_layout) live on - # the offload SUB-connector, not the composite. Passing the raw - # composite here read a zero chunk grid. Unwrap the sub the same way - # `state_tier_capability` does, so the gate's chunk size and the tier - # the capability check builds stay consistent. - if isinstance(kvcfg, dict) and kvcfg.get("kv_connector") == "multi": - sub, _why = _offload_subconfig(kvcfg) - if sub is not None: - kvcfg = sub - return build_lmcache_config(kvcfg) - - def state_tier_chunk_tokens(config) -> int: """The LMCache chunk size in tokens for the joint KV leg, or 0. @@ -639,7 +616,19 @@ def state_tier_chunk_tokens(config) -> int: disable the joint KV load. """ try: - return int(_lmcache_engine_config(config).chunk_size) + from atom.kv_transfer.offload.config import build_lmcache_config + + kvcfg = getattr(config, "kv_transfer_config", None) or {} + # Under `multi` the lmcache.* keys (chunk_size, offload_layout) live on + # the offload SUB-connector, not the composite. Passing the raw + # composite here read a zero chunk grid. Unwrap the sub the same way + # `state_tier_capability` does, so the gate's chunk size and the tier + # the capability check builds stay consistent. + if isinstance(kvcfg, dict) and kvcfg.get("kv_connector") == "multi": + sub, _why = _offload_subconfig(kvcfg) + if sub is not None: + kvcfg = sub + return int(build_lmcache_config(kvcfg).chunk_size) except Exception: # Blind on purpose: this runs at model load, the import reaches a # third-party package that may be absent entirely, and the only @@ -658,7 +647,15 @@ def state_tier_verified_for_s(config) -> float: when it cannot be read. Like `state_tier_chunk_tokens`, call it only when a tier is hosted.""" try: - pool_gib = float(_lmcache_engine_config(config).max_local_cpu_size) + from atom.kv_transfer.offload.config import build_lmcache_config + + kvcfg = getattr(config, "kv_transfer_config", None) or {} + # The offload sub-connector's keys, as in `state_tier_chunk_tokens`. + if isinstance(kvcfg, dict) and kvcfg.get("kv_connector") == "multi": + sub, _why = _offload_subconfig(kvcfg) + if sub is not None: + kvcfg = sub + pool_gib = float(build_lmcache_config(kvcfg).max_local_cpu_size) except Exception: # noqa: BLE001 - `state_tier_chunk_tokens` logged why logger.warning( "state offload: LMCache's CPU pool size is unknown; a confirmed " diff --git a/tests/conftest.py b/tests/conftest.py index 22ca235a7d..209a2a47f6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -31,7 +31,6 @@ from atom.model_engine.block_manager import BlockManager from atom.model_engine.scheduler import Scheduler from atom.model_engine.sequence import Sequence -from atom.model_engine.state_offload import _TIER_VERIFIED_MAX_S from atom.sampling_params import SamplingParams from atom.utils.region_graph import ReqGraph @@ -248,18 +247,6 @@ def keep_envs_lazy(): delattr(envs, name) -#: Past every confirmation window `tier_verified_for_s` can give. -OFFLOAD_STALE_S = _TIER_VERIFIED_MAX_S + 1 - - -@pytest.fixture -def offload_clock(monkeypatch): - """`state_offload`'s monotonic clock, advanced by hand.""" - now = [1000.0] - monkeypatch.setattr("atom.model_engine.state_offload.monotonic", lambda: now[0]) - return now - - @pytest.fixture def seq_factory(): """Factory for creating Sequence objects with sensible defaults.""" diff --git a/tests/test_block_manager.py b/tests/test_block_manager.py index 7533a9a4cb..ea3853f6db 100644 --- a/tests/test_block_manager.py +++ b/tests/test_block_manager.py @@ -1381,9 +1381,8 @@ def test_the_block_manager_sizes_the_window_from_the_lmcache_pool( self, monkeypatch ): monkeypatch.setattr( - so, - "_lmcache_engine_config", - lambda config: SimpleNamespace(chunk_size=256, max_local_cpu_size=100), + "atom.kv_transfer.offload.config.build_lmcache_config", + lambda kvcfg: SimpleNamespace(chunk_size=256, max_local_cpu_size=100), ) bm = _bm_with_state_tier(monkeypatch, _OFFLOAD_KVC) assert bm._joint_chunk_tokens == 256 diff --git a/tests/test_lmcache_offload_connector.py b/tests/test_lmcache_offload_connector.py index 59c181f5a4..e607fb8f39 100644 --- a/tests/test_lmcache_offload_connector.py +++ b/tests/test_lmcache_offload_connector.py @@ -6162,22 +6162,11 @@ def test_no_more_saves_go_out_than_the_pool_can_afford_to_pin(): # ── the state key: build safety and namespace separation ────────────────── -class _LRUPolicy: - """LMCache's `LRUCachePolicy.update_on_hit`.""" - - def update_on_hit(self, key, cache) -> None: - cache.move_to_end(key) - - class _FakeStorage: - """The `storage_manager` calls the codec makes, over a plain dict, with - the LocalCPU LRU internals `bind_storage_manager` requires.""" + """The two `storage_manager` calls the codec makes, over a plain dict.""" def __init__(self) -> None: - self.objects: OrderedDict = OrderedDict() - self.local_cpu_backend = SimpleNamespace( - cpu_lock=threading.Lock(), hot_cache=self.objects, cache_policy=_LRUPolicy() - ) + self.objects: dict = {} def batched_put(self, keys, objs) -> None: for k, o in zip(keys, objs, strict=True): @@ -7185,9 +7174,6 @@ def test_the_scheduler_half_keeps_the_operation_whole(self): assert failed == set() -# ── kimi_k3: loads must not queue behind a backlog of stores ─────────────── - - class _Obj: """A MemoryObj's reference count, and nothing else.""" @@ -7209,6 +7195,13 @@ def get_size(self) -> int: return self.size +class _LRUPolicy: + """LMCache's `LRUCachePolicy.update_on_hit`.""" + + def update_on_hit(self, key, cache) -> None: + cache.move_to_end(key) + + class _LocalCPU: """`storage_manager` over LocalCPU's rules. @@ -7342,14 +7335,6 @@ def test_a_resume_refreshes_the_image(self, lmcache_key): assert objs[0].refs == 1, "the load's reference is discharged" assert pool.evict_oldest() == "kv-chunk" - def test_a_missed_get_leaves_the_cap_lru(self, lmcache_key): - pool = _LocalCPU() - codec, _ = _lru_codec(pool, max_entries=2) - assert codec.put(1, []) - pool.evict_oldest() - assert not codec.get(1, 0) - assert codec._cache_keys == OrderedDict() - def test_a_resident_image_is_refreshed_not_stored_twice(self, lmcache_key): pool = _LocalCPU() codec, objs = _lru_codec(pool, max_entries=2) @@ -7607,6 +7592,9 @@ def test_both_touch_reports_reach_quorum_across_tp(self, monkeypatch): } +# ── kimi_k3: loads must not queue behind a backlog of stores ─────────────── + + class TestStateLoadsAndStoresRunInSeparateLanes: """A load is on the TTFT critical path and a store is not, but one serial executor made that ordering unenforceable: a load submitted in a later step diff --git a/tests/test_state_checkpoint.py b/tests/test_state_checkpoint.py index 84794c1def..ced430f71e 100644 --- a/tests/test_state_checkpoint.py +++ b/tests/test_state_checkpoint.py @@ -17,7 +17,6 @@ from types import SimpleNamespace import pytest -from conftest import OFFLOAD_STALE_S as STALE from conftest import MockConfig from atom.model_engine.block_manager import ( @@ -35,7 +34,7 @@ from atom.model_engine.scheduler import ScheduledBatchOutput, Scheduler from atom.model_engine.sequence import Sequence, SequenceType from atom.model_engine.state_cache import StateCache -from atom.model_engine.state_offload import StateOffloadIndex +from atom.model_engine.state_offload import _TIER_VERIFIED_MAX_S, StateOffloadIndex from atom.model_engine.state_pool import StateSlotPool from atom.model_engine.state_runtime import ( StateRuntime, @@ -4567,6 +4566,18 @@ def test_a_cold_start_without_a_boundary_is_untouched(self): # ── state tier hit refresh: the admission hold and the voters ────────────── +#: Past every confirmation window `tier_verified_for_s` can give. +STALE = _TIER_VERIFIED_MAX_S + 1 + + +@pytest.fixture +def offload_clock(monkeypatch): + """`state_offload`'s monotonic clock, advanced by hand.""" + now = [1000.0] + monkeypatch.setattr("atom.model_engine.state_offload.monotonic", lambda: now[0]) + return now + + def _paged_bm(): bm = make_block_manager( ckpt_config(num_kvcache_blocks=160, pool_entries={"state": 8}), @@ -4747,19 +4758,7 @@ def test_asking_whether_the_tier_can_serve_registers_nothing(self, offload_clock class TestResumeHash: - """`allocate` and the admission hold pick the state to resume from with - one helper, `_state_resume_hash`.""" - - def test_the_joint_boundary_wins_then_a_replay_then_the_hit(self, monkeypatch): - bm = _paged_bm() - seq = _seq() - assert bm._state_resume_hash(seq, 77, 2) == 77 - assert bm._state_resume_hash(seq, -1, 2) == 2 - assert bm._state_resume_hash(seq, -1, -1) == -1 - monkeypatch.setattr(bm, "state_replay_tokens", 8) - seq.replays_hit = True - assert bm._state_resume_hash(seq, -1, 2) == -1, "a replay restores none" - assert bm._state_resume_hash(seq, 77, 2) == 77 + """The admission hold judges the hash `allocate` then restores.""" def test_the_hold_judges_the_hash_allocate_then_loads( self, offload_clock, monkeypatch diff --git a/tests/test_state_offload_index.py b/tests/test_state_offload_index.py index 8c5cecc224..462da72b45 100644 --- a/tests/test_state_offload_index.py +++ b/tests/test_state_offload_index.py @@ -1,12 +1,9 @@ # SPDX-License-Identifier: MIT -"""StateOffloadIndex must stop advertising images the worker codec dropped, -and keeps state images as warm in LMCache's LRU as the KV they belong to -(hit touches, their verdicts, and the admission hold).""" +"""StateOffloadIndex must stop advertising images the worker codec dropped.""" from types import SimpleNamespace import pytest -from conftest import OFFLOAD_STALE_S as STALE from atom.model_engine import state_offload as so from atom.model_engine.state_offload import ( @@ -66,6 +63,18 @@ def test_state_tier_cpu_bytes_reads_gib(monkeypatch): # ── the index: votes queue touches, verdicts keep it in step ─────────────── +#: Past every confirmation window `tier_verified_for_s` can give. +STALE = so._TIER_VERIFIED_MAX_S + 1 + + +@pytest.fixture +def offload_clock(monkeypatch): + """`state_offload`'s monotonic clock, advanced by hand.""" + now = [1000.0] + monkeypatch.setattr("atom.model_engine.state_offload.monotonic", lambda: now[0]) + return now + + def _index(*hashes, verified_for_s=so._TIER_VERIFIED_MIN_S, **kw): idx = StateOffloadIndex(verified_for_s=verified_for_s, **kw) for h in hashes: From 3ccc2e9fe0f01a1f7d84a415f76f2fa2d17b89c2 Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Mon, 2 Nov 2026 10:01:33 +0000 Subject: [PATCH 16/18] state offload: keep the shared resume-hash helper, drop the budget warning - block_manager.py: _state_resume_hash is back, shared by allocate and the admission hold as asked in review, so the two cannot drift apart. record_allocation keeps its original shape. - Drop the startup state-cap / pool-size warning; it belongs to the removed reservation, not to the hit refresh. - tests: drop the four codec tests that pass on main unchanged (they test _put, which this PR does not change) and the fake LocalCPU knobs only they used. - scheduler.py: say the hold is base-scheduler only. --- atom/kv_transfer/offload/README.md | 3 +- .../offload/hybrid/kimi_k3/connector.py | 11 +-- atom/model_engine/block_manager.py | 74 +++++++++++-------- atom/model_engine/scheduler.py | 3 +- atom/model_engine/state_offload.py | 24 ------ tests/test_lmcache_offload_connector.py | 39 +--------- tests/test_state_checkpoint.py | 14 +++- tests/test_state_offload_index.py | 21 +----- 8 files changed, 65 insertions(+), 124 deletions(-) diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index d4187713d1..356ba764a8 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -131,8 +131,7 @@ 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`. The worker warns when -`OFFLOAD_STATE_CPU_SIZE` is more than half of the pool. +the connector statistic `state_touches_unrefreshed`. ## LMCache Multiprocess (`lmcache_mp`) diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py index 190b446664..5c2816b55d 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py @@ -40,10 +40,7 @@ _JointPark, ) from atom.kv_transfer.offload.metadata import LMCacheOffloadMetadata -from atom.model_engine.state_offload import ( - state_tier_budget_warning, - state_tier_cpu_bytes, -) +from atom.model_engine.state_offload import state_tier_cpu_bytes from atom.utils import envs logger = logging.getLogger("atom") @@ -262,12 +259,6 @@ def _build_state_tier(self, transfer_tensors) -> None: # Share the allocator, but cap state retention separately so old # checkpoints cannot consume nearly the whole paged-KV CPU pool. codec.bind_storage_manager(self._engine.storage_manager) - pool_gib = float(getattr(cfg, "max_local_cpu_size", 0) or 0) - budget_warning = state_tier_budget_warning( - state_tier_cpu_bytes() / (1 << 30), pool_gib - ) - if budget_warning: - logger.warning("kimi_k3 offload: %s", budget_warning) # No index here: StateOffloadIndex lives in the engine process; both # directions report and the engine applies. self._state_tier = StateOffloadTier(codec) diff --git a/atom/model_engine/block_manager.py b/atom/model_engine/block_manager.py index 7ddc310cc7..4ca0fec712 100644 --- a/atom/model_engine/block_manager.py +++ b/atom/model_engine/block_manager.py @@ -1577,28 +1577,41 @@ def _resume_hash_off_hbm( block_hashes: list[int], decision: JointBoundaryDecision, ) -> int: - """The hash `allocate` would restore `seq`'s state from, when HBM does - not hold it; -1 for no state, an HBM-held one, or a replay. - - Mirrors `allocate`'s choice: the joint boundary when there is one (above - the HBM hit by construction), else the HBM hit's last block unless a - replay rebuilds the state instead. - """ + """The `_state_resume_hash` `allocate` would restore from, when HBM + does not hold it; -1 for no state, an HBM-held one, or a replay.""" if not seq.has_per_req_cache: return -1 - if decision.skip_reason is None: - return -1 if decision.state_from_hbm else decision.boundary_hash - if self.tier_resume_replays(seq) or not ( - 0 < num_cached_blocks <= len(block_hashes) - ): + joint = decision.skip_reason is None + if joint and decision.state_from_hbm: return -1 - h = block_hashes[num_cached_blocks - 1] - in_hbm = ( - self.paged_state_checkpoints.contains(h) - if self.paged_state_checkpoints is not None - else self.state.lookup(h) >= 0 + hit_hash = ( + block_hashes[num_cached_blocks - 1] + if 0 < num_cached_blocks <= len(block_hashes) + else -1 + ) + h = self._state_resume_hash( + seq, decision.boundary_hash if joint else -1, hit_hash ) - return -1 if in_hbm else h + if h != -1 and not joint: + in_hbm = ( + self.paged_state_checkpoints.contains(h) + if self.paged_state_checkpoints is not None + else self.state.lookup(h) >= 0 + ) + if in_hbm: + return -1 + return h + + def _state_resume_hash(self, seq: Sequence, joint_hash: int, hit_hash: int) -> int: + """The hash whose state `seq` resumes from, or -1 for none. + + The joint boundary when there is one (above the HBM hit by + construction), else the HBM hit's last block. A replay rebuilds the + state instead of restoring it, so it resumes from none. + """ + if joint_hash != -1: + return joint_hash + return -1 if self.tier_resume_replays(seq) else hit_hash def retreat_joint_to_hbm(self, seq: Sequence) -> bool: """Drop a failed unaligned joint and resume at the HBM hit it lowered. @@ -1795,13 +1808,16 @@ def allocate(self, seq: Sequence, num_cached_blocks: int = 0) -> bool: # last one actually claimed, whichever way the loop exits. if i < num_cached_blocks: hit_hash = h - if self.state_replay_tokens and seq.replays_hit: - # A replay starts from a fresh state slot, never a checkpoint: it - # rebuilds the state below whatever the claim above reached. - hit_hash = -1 - if num_cached_blocks < requested_blocks: - # A KV-only load past the hit would start above a gap. - seq.replay_kv_load = False + # A replay starts from a fresh state slot, never a checkpoint: it + # rebuilds the state below whatever the claim above reached + # (`_state_resume_hash` gives it none). A KV-only load past the hit + # would start above a gap. + if ( + self.state_replay_tokens + and seq.replays_hit + and num_cached_blocks < requested_blocks + ): + seq.replay_kv_load = False # Pin the restore before fresh blocks can evict its checkpoint. state_holds = True if seq.has_per_req_cache and self.paged_state_checkpoints is not None: @@ -1817,9 +1833,9 @@ def allocate(self, seq: Sequence, num_cached_blocks: int = 0) -> bool: # joint boundary and `_claim_after_load` raised `num_cached_tokens` # to it, so the forward resumed over a prefix the state does not # cover. Silent wrong output, no exception. - joint_hash = seq.offload_joint.boundary_hash state_holds = self._attach_state_slots( - seq, joint_hash if joint_hash != -1 else hit_hash + seq, + self._state_resume_hash(seq, seq.offload_joint.boundary_hash, hit_hash), ) for _ in range(len(seq.block_table), self.num_pool_blocks(len(seq))): seq.block_table.append(self._fresh_block()) @@ -1836,9 +1852,9 @@ def allocate(self, seq: Sequence, num_cached_blocks: int = 0) -> bool: # above the HBM hit by construction. `num_cached_tokens` stays at # the HBM prefix until the KV leg lands, so the forward covers # `[hbm, num_prompt)` if anything goes wrong from here. - joint_hash = seq.offload_joint.boundary_hash state_holds = self._attach_state_slots( - seq, joint_hash if joint_hash != -1 else hit_hash + seq, + self._state_resume_hash(seq, seq.offload_joint.boundary_hash, hit_hash), ) if seq.has_per_req_cache: seq._state_initialized_after_alloc = False diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index c25e633a64..0e692c9fc4 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -2384,7 +2384,8 @@ def _schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]] | None: ): # Held a pass for the state offload's verdict on the image it # would load; a gone image is forgotten and the retry plans - # without it. + # without it. Base scheduler only: DecodeScheduler never + # publishes touches, so it never holds. prefix_waiters.append(seq) continue if not self.block_manager.allocate(seq, num_cached_blocks): diff --git a/atom/model_engine/state_offload.py b/atom/model_engine/state_offload.py index 2c50e62ae3..19b6495282 100644 --- a/atom/model_engine/state_offload.py +++ b/atom/model_engine/state_offload.py @@ -574,30 +574,6 @@ def state_tier_cpu_bytes() -> int: return int(float(os.environ.get("OFFLOAD_STATE_CPU_SIZE", "32")) * (1 << 30)) -def state_tier_budget_warning(state_gib: float, pool_gib: float) -> str | None: - """Why the state cap and the LMCache pool it shares do not fit together - (e.g. `LMCACHE_MAX_LOCAL_CPU_SIZE=32` with the 32 GiB default cap), or - None.""" - if pool_gib <= 0: - return None - sizes = ( - f"OFFLOAD_STATE_CPU_SIZE={state_gib:.1f} GiB, " - f"LMCACHE_MAX_LOCAL_CPU_SIZE={pool_gib:.1f} GiB" - ) - if state_gib >= pool_gib: - return ( - f"{sizes}: the state cap is not below the pool, so it bounds " - "nothing and state images can take the whole pool from paged KV. " - "Lower OFFLOAD_STATE_CPU_SIZE or raise LMCACHE_MAX_LOCAL_CPU_SIZE." - ) - if 2 * state_gib > pool_gib: - return ( - f"{sizes}: state images may take more than half of the pool, " - "leaving paged KV the smaller share under pressure." - ) - return None - - def state_tier_chunk_tokens(config) -> int: """The LMCache chunk size in tokens for the joint KV leg, or 0. diff --git a/tests/test_lmcache_offload_connector.py b/tests/test_lmcache_offload_connector.py index e607fb8f39..db1b484274 100644 --- a/tests/test_lmcache_offload_connector.py +++ b/tests/test_lmcache_offload_connector.py @@ -7212,24 +7212,20 @@ class _LocalCPU: anything another holder still references. """ - def __init__(self, refuse_all: bool = False) -> None: + def __init__(self) -> None: self.hot: OrderedDict = OrderedDict() # A second backend (LocalDisk) for presence only: `contains` answers # from it, `get` does not read it (no test loads a disk-only image). self.disk: dict = {} - self.refuse_all = refuse_all - self.fail_put = False self.puts = 0 self.local_cpu_backend = SimpleNamespace( cpu_lock=threading.Lock(), hot_cache=self.hot, cache_policy=_LRUPolicy() ) def batched_put(self, keys, objs) -> None: - if self.fail_put: - raise RuntimeError("put failed") self.puts += 1 for k, o in zip(keys, objs, strict=True): - if not self.refuse_all and k not in self.hot: + if k not in self.hot: o.ref_count_up() self.hot[k] = o o.ref_count_down() @@ -7299,15 +7295,6 @@ def _kv(pool, name): class TestCodecRefresh: - def test_an_untouched_image_is_the_first_thing_pressure_evicts(self, lmcache_key): - """The bug: KV written after the image is younger, however often the - prefix hits.""" - pool = _LocalCPU() - codec, _ = _lru_codec(pool, max_entries=None) - assert codec.put(1, []) - _kv(pool, "kv-chunk") - assert pool.evict_oldest() == codec.key(1) - def test_a_touch_moves_the_image_past_older_kv(self, lmcache_key): pool = _LocalCPU() codec, _ = _lru_codec(pool, max_entries=None) @@ -7367,28 +7354,6 @@ def test_the_cap_retires_the_least_recently_used_image(self, lmcache_key): assert objs[1].refs == 0, "a retired image goes back to the allocator" assert list(codec._cache_keys) == [codec.key(1), codec.key(3)] - def test_images_hold_no_extra_reference(self, lmcache_key): - """No reservation: KV pressure may evict an image like any chunk.""" - pool = _LocalCPU() - codec, objs = _lru_codec(pool, max_entries=2) - assert codec.put(1, []) - assert objs[0].refs == 1 - assert pool.evict_oldest() == codec.key(1) - - def test_a_refused_put_keeps_no_reference(self, lmcache_key): - pool = _LocalCPU(refuse_all=True) - codec, objs = _lru_codec(pool, max_entries=2) - assert codec.put(1, []) - assert objs[0].refs == 0, "holding a refused object would strand its block" - - def test_a_failed_put_discharges_its_reference(self, lmcache_key): - pool = _LocalCPU() - pool.fail_put = True - codec, objs = _lru_codec(pool, max_entries=2) - with pytest.raises(RuntimeError): - codec.put(1, []) - assert objs[0].refs == 0 - @pytest.mark.parametrize( "missing", [ diff --git a/tests/test_state_checkpoint.py b/tests/test_state_checkpoint.py index ced430f71e..672b62d0f5 100644 --- a/tests/test_state_checkpoint.py +++ b/tests/test_state_checkpoint.py @@ -4758,7 +4758,19 @@ def test_asking_whether_the_tier_can_serve_registers_nothing(self, offload_clock class TestResumeHash: - """The admission hold judges the hash `allocate` then restores.""" + """`allocate` and the admission hold pick the state to resume from with + one helper, `_state_resume_hash`.""" + + def test_the_joint_boundary_wins_then_a_replay_then_the_hit(self, monkeypatch): + bm = _paged_bm() + seq = _seq() + assert bm._state_resume_hash(seq, 77, 2) == 77 + assert bm._state_resume_hash(seq, -1, 2) == 2 + assert bm._state_resume_hash(seq, -1, -1) == -1 + monkeypatch.setattr(bm, "state_replay_tokens", 8) + seq.replays_hit = True + assert bm._state_resume_hash(seq, -1, 2) == -1, "a replay restores none" + assert bm._state_resume_hash(seq, 77, 2) == 77 def test_the_hold_judges_the_hash_allocate_then_loads( self, offload_clock, monkeypatch diff --git a/tests/test_state_offload_index.py b/tests/test_state_offload_index.py index 462da72b45..e9be6c39bf 100644 --- a/tests/test_state_offload_index.py +++ b/tests/test_state_offload_index.py @@ -6,11 +6,7 @@ import pytest from atom.model_engine import state_offload as so -from atom.model_engine.state_offload import ( - StateOffloadIndex, - state_tier_budget_warning, - state_tier_cpu_bytes, -) +from atom.model_engine.state_offload import StateOffloadIndex, state_tier_cpu_bytes def _stored(idx, hashes): @@ -310,18 +306,3 @@ def test_verdicts_are_collected(self, offload_clock): SimpleNamespace(take_state_touch_verdicts=lambda: {7: False, 8: True}) ) assert not idx.could_serve(7) and idx.could_serve(8) - - -@pytest.mark.parametrize( - "state_gib, pool_gib, says", - [ - (59, 208, None), - (32, 32, "bounds nothing"), # #2562's accuracy run: the default cap - (40, 64, "more than half"), - (32, 64, None), - (32, 0, None), - ], -) -def test_the_state_cap_is_judged_against_the_pool(state_gib, pool_gib, says): - msg = state_tier_budget_warning(state_gib, pool_gib) - assert msg is None if says is None else says in msg From bdef7efb4f165b85ea036cadc253dcac5545bcc1 Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Thu, 5 Nov 2026 02:53:12 +0000 Subject: [PATCH 17/18] state offload: correct the hold and touch-verdict comments, tidy test imports - scheduler.py: DecodeScheduler does publish touches when idle; it never holds because allocate_waiting never passes defer_unverified. - kimi_k3/connector.py: say that a GONE verdict counts as refreshed. - state_offload.py: name loads_deferred_for_verify in the stats comment; an expired hold means verdicts are lost or late (e.g. a stalled step), not necessarily lost. - tests: import the touch channels with the other kimi_k3 names at the top instead of inside test functions. --- .../offload/hybrid/kimi_k3/connector.py | 1 + atom/model_engine/scheduler.py | 5 +++-- atom/model_engine/state_offload.py | 10 ++++++---- tests/test_lmcache_offload_connector.py | 17 ++--------------- 4 files changed, 12 insertions(+), 21 deletions(-) diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py index 5c2816b55d..416df2ae1e 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py @@ -128,6 +128,7 @@ def __init__(self, config) -> None: # 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 diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index ae9b8204bb..5f9d685d2b 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -2435,8 +2435,9 @@ def _schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]] | None: ): # Held a pass for the state offload's verdict on the image it # would load; a gone image is forgotten and the retry plans - # without it. Base scheduler only: DecodeScheduler never - # publishes touches, so it never holds. + # without it. Base scheduler only: DecodeScheduler's + # `allocate_waiting` never passes `defer_unverified`, so it + # never holds. prefix_waiters.append(seq) continue if not self.block_manager.allocate(seq, num_cached_blocks): diff --git a/atom/model_engine/state_offload.py b/atom/model_engine/state_offload.py index 19b6495282..f24db354fb 100644 --- a/atom/model_engine/state_offload.py +++ b/atom/model_engine/state_offload.py @@ -280,8 +280,8 @@ def hold_admission(self, h: int) -> bool: logger.warning( "state offload: no touch verdict for a CPU state image " "within OFFLOAD_LOOKUP_DEFER_S=%.1fs; admitting without it. " - "Touch verdicts are being lost (counted as " - "state_offload_verify_holds_expired).", + "Touch verdicts are lost or late, e.g. a stalled step " + "(counted as state_offload_verify_holds_expired).", defer_s, ) return False @@ -385,8 +385,10 @@ def stats(self) -> dict[str, int]: "hashes_evicted": self.hashes_evicted, # Hit refresh. `touch_misses`: images a touch found gone, which # left the index (one a load was already offered on also counts in - # `loads_failed`). `verify_holds_expired`: holds that got no - # verdict within `defer_s`; non-zero means verdicts are being lost. + # `loads_failed`). `loads_deferred_for_verify`: admission passes + # held. `verify_holds_expired`: holds that got no verdict within + # `defer_s`; non-zero means verdicts are lost or late (e.g. a + # stalled step). "touches_sent": self.touches_sent, "touch_misses": self.touch_misses, "loads_deferred_for_verify": self.loads_deferred_for_verify, diff --git a/tests/test_lmcache_offload_connector.py b/tests/test_lmcache_offload_connector.py index db1b484274..0b43aa3eda 100644 --- a/tests/test_lmcache_offload_connector.py +++ b/tests/test_lmcache_offload_connector.py @@ -75,6 +75,8 @@ ) from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( STATE_INDEX_CHANNEL, + STATE_TOUCH_CHANNEL, + STATE_TOUCH_REFRESH_CHANNEL, KimiK3OffloadConnector, KimiK3OffloadScheduler, save_stall_seconds, @@ -7427,10 +7429,6 @@ def test_a_disk_only_image_is_stored_back_into_the_cpu_tier(self, lmcache_key): def _touch_worker(tier): - from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( - KimiK3OffloadConnector, - ) - c = KimiK3OffloadConnector.__new__(KimiK3OffloadConnector) c._state_tier = tier c._touch_verdicts = {} @@ -7498,10 +7496,6 @@ def test_each_touch_is_a_distinct_identity_drained_once(self, monkeypatch): assert s.build_connector_meta().state_touches == [] def test_one_rank_missing_the_image_makes_it_gone(self, monkeypatch): - from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( - STATE_TOUCH_CHANNEL, - ) - s = self._scheduler(monkeypatch) for op, ok in ((StateTouchId(5, 0), True), (StateTouchId(5, 1), False)): handled = s.connector_completion( @@ -7518,13 +7512,6 @@ def test_both_touch_reports_reach_quorum_across_tp(self, monkeypatch): """Each rank files a residency and a refresh report per touch; one rank that only found the image present counts the touch unrefreshed without making it gone.""" - from atom.kv_transfer.disaggregation.aggregator import KVOutputAggregator - from atom.kv_transfer.disaggregation.types import KVConnectorOutput - from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( - STATE_TOUCH_CHANNEL, - STATE_TOUCH_REFRESH_CHANNEL, - ) - touch = StateTouchId(5, 0) outputs = [] for result in (TouchResult.REFRESHED, TouchResult.PRESENT): From 8d407822219baad5f72ddc3b9b3ecd32f453d160 Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Fri, 6 Nov 2026 03:59:41 +0000 Subject: [PATCH 18/18] Simplify the state touch path and its tests - K3 worker keeps each touch's TouchResult (a plain Enum, compared by identity) and turns it into the two completions in `_emit_touch_verdicts`, instead of a (resident, refreshed) tuple built up front. With no tier, every touch is GONE before the loop. - `state_tier_chunk_tokens` and `state_tier_verified_for_s` share one reader, `_lmcache_engine_config`; the import, the `multi` unwrap and the build move into it unchanged. - Conversions done once: touch ids are ints where `enqueue_state_touches` makes them; the later int()/bool()/float() calls are gone. `-inf` instead of a -1e18 sentinel. - Comments state constraints only; the measurement behind the window constants is in the PR description, and the offload README no longer repeats them. - Tests check emitted completions and pool contents rather than private state; `offload_clock` lives in conftest; STALE and the window floor come from `tier_verified_for_s`; no function-local imports. - New tests for the wiring the existing ones missed: a touch-only step is work, `start_load_kv` touches, the scheduler applies verdicts, and both scheduler shells (plain and multi) forward touches and verdicts. No behaviour change: same values and branches. Checked with 34 mutants of the touch path (all killed), the CPU suite (same 3 failures as main), GPU unit tests (1863 passed; the 2 engram failures are on main too) and Kimi-K3 TP2+EP2 GSM8K 0.978 on the previous head e9481f9bd. Limit: no GSM8K rerun on this commit, and that run sent no touches. --- atom/kv_transfer/offload/README.md | 2 +- .../offload/hybrid/kimi_k3/connector.py | 53 ++++++----- .../offload/hybrid/kimi_k3/state_object.py | 14 +-- atom/model_engine/state_offload.py | 61 ++++++------- tests/conftest.py | 8 ++ tests/test_block_manager.py | 2 +- tests/test_lmcache_offload_connector.py | 90 ++++++++++++------- tests/test_offload_shell.py | 40 +++++++++ tests/test_state_checkpoint.py | 32 ++++--- tests/test_state_offload_index.py | 38 ++------ 10 files changed, 201 insertions(+), 139 deletions(-) diff --git a/atom/kv_transfer/offload/README.md b/atom/kv_transfer/offload/README.md index 23994a23fa..8ae03fb3a1 100644 --- a/atom/kv_transfer/offload/README.md +++ b/atom/kv_transfer/offload/README.md @@ -128,7 +128,7 @@ 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 +(sized from `LMCACHE_MAX_LOCAL_CPU_SIZE`) 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`. diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py index 416df2ae1e..336d3ecb84 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/connector.py @@ -127,9 +127,8 @@ 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]] = {} + # Touch verdicts since the last `get_finished`. + self._touch_verdicts: dict[StateTouchId, TouchResult] = {} # Only the first touch whose codec call raised is logged. self._warned_touch_failure = False @@ -311,12 +310,12 @@ def _touch_states(self, metadata) -> None: if not touches: return tier = self._state_tier + if tier is None: + self._touch_verdicts.update(dict.fromkeys(touches, TouchResult.GONE)) + return 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)) + result = tier.codec.touch(touch_id.prefix_hash) except Exception: # Says nothing about the bytes: kept, not refreshed. One # traceback, since a drifted LMCache raises on every touch. @@ -330,10 +329,7 @@ def _touch_states(self, metadata) -> None: exc_info=True, ) result = TouchResult.PRESENT - self._touch_verdicts[touch_id] = ( - bool(result), - result is not TouchResult.PRESENT, - ) + self._touch_verdicts[touch_id] = result def _arm_joint_loads(self, metadata) -> None: """Hold a request owning both legs until both report. @@ -516,14 +512,7 @@ def get_finished(self): 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 = {} + self._emit_touch_verdicts(out) 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, @@ -610,6 +599,24 @@ def _emit_load_dispositions(self, out) -> None: ) ) + def _emit_touch_verdicts(self, out) -> None: + """Two completions per touch since the last call: still resident, and + refreshed (GONE counts as refreshed: nothing was left unrefreshed).""" + for touch_id, result in self._touch_verdicts.items(): + out.connector_completions.add( + ConnectorCompletion( + STATE_TOUCH_CHANNEL, touch_id, result is not TouchResult.GONE + ) + ) + out.connector_completions.add( + ConnectorCompletion( + STATE_TOUCH_REFRESH_CHANNEL, + touch_id, + result is not TouchResult.PRESENT, + ) + ) + self._touch_verdicts = {} + def _reclaim_stale_parks(self) -> None: """Evict joint parks whose report never came (finding #3). @@ -948,10 +955,10 @@ def connector_completion(self, completion) -> bool | None: 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) + h = completion.operation_id.prefix_hash + self._state_touch_verdicts[h] = ( + self._state_touch_verdicts.get(h, True) and completion.succeeded + ) # Handled, but not a terminal save: `process_completions` would # file the touch id under `finished_saving`. return None diff --git a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py index 4ec6da918a..9f91a842c7 100644 --- a/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py +++ b/atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py @@ -22,14 +22,14 @@ logger = logging.getLogger("atom") -class TouchResult(enum.IntEnum): - """What a touch found. Truthy when the image can still be read.""" +class TouchResult(enum.Enum): + """What a touch found.""" - GONE = 0 + GONE = enum.auto() #: Held only by another backend: present, not refreshed. `get` still reads #: it; LMCache's `StorageManager.get` writes it back into LocalCPU. - PRESENT = 1 - REFRESHED = 2 + PRESENT = enum.auto() + REFRESHED = enum.auto() def _local_cpu_lru(storage_manager): @@ -188,7 +188,7 @@ def put(self, h: int, unit_ids, on_source_released=None) -> bool: def touch(self, h: int) -> TouchResult: """Refresh `h`'s image in LMCache's LRU, as a KV lookup does for its - chunks; falsy (`GONE`) when it is gone. A state hit is answered by + chunks; `GONE` when it is gone. A state hit is answered by `StateOffloadIndex`, so otherwise LMCache never sees the image used. """ if self._storage is None: @@ -212,7 +212,7 @@ def touch(self, h: int) -> TouchResult: def _touch_key(self, key) -> TouchResult: result = self._refresh(key) with self._cache_lock: - if result: + if result is not TouchResult.GONE: if key in self._cache_keys: self._cache_keys.move_to_end(key) else: diff --git a/atom/model_engine/state_offload.py b/atom/model_engine/state_offload.py index f24db354fb..b94ad3cb5e 100644 --- a/atom/model_engine/state_offload.py +++ b/atom/model_engine/state_offload.py @@ -15,6 +15,7 @@ import os from collections import OrderedDict from dataclasses import dataclass, replace +from math import inf from time import monotonic from atom.utils import envs @@ -22,9 +23,9 @@ logger = logging.getLogger("atom") #: Seconds a confirmed image is trusted per GiB of LMCache pool: a quarter of -#: the pool's turnover at 1 GiB/s, four times the fill rate AgentX c64 measured. +#: the time the pool takes to turn over at 1 GiB/s of new stores. _TIER_VERIFIED_S_PER_GIB = 0.25 -#: Well above a touch round trip / the largest window measured (AgentX c64). +#: Well above a touch round trip. _TIER_VERIFIED_MIN_S = 5.0 _TIER_VERIFIED_MAX_S = 60.0 @@ -106,10 +107,10 @@ def __init__( self.stores_untrusted = 0 # Off for indexes whose connector carries no touches (the vLLM plugin). self.refresh_hits = verified_for_s is not None - self.verified_for_s = float(verified_for_s or 0.0) + self.verified_for_s = verified_for_s or 0.0 # Read once: `hold_admission` runs on the admission path. Not read # without the refresh, which never holds. - self.defer_s = float(envs.OFFLOAD_LOOKUP_DEFER_S) if self.refresh_hits else 0.0 + self.defer_s = envs.OFFLOAD_LOOKUP_DEFER_S if self.refresh_hits else 0.0 self._touch_wanted: set[int] = set() # hash -> monotonic of the last touch sent / the last evidence the # image was resident. @@ -200,7 +201,7 @@ def take_touches(self) -> list[int]: now = monotonic() out = [] for h in self._touch_wanted: - if now - self._touch_sent_at.get(h, -1e18) < _TIER_TOUCH_INTERVAL_S: + if now - self._touch_sent_at.get(h, -inf) < _TIER_TOUCH_INTERVAL_S: continue self._touch_sent_at[h] = now out.append(h) @@ -210,12 +211,11 @@ def take_touches(self) -> list[int]: def apply_touch_verdict(self, h: int, resident: bool) -> None: """A touch came back; gone drops `h` before a request parks on it.""" - h = int(h) if h not in self.hashes: return if not resident: # Stale: `h` was stored (again) after this touch looked. - if self._verified_at.get(h, -1e18) > self._touch_sent_at.get(h, -1e18): + if self._verified_at.get(h, -inf) > self._touch_sent_at.get(h, -inf): return self.touch_misses += 1 self.forget(h) @@ -257,10 +257,7 @@ def hold_admission(self, h: int) -> bool: outside HBM) waits a pass for its touch verdict: `h` not confirmed within `verified_for_s`. If the verdict is gone, `h` has left the index and the retry plans a lower boundary instead of parking on a miss. - - Bounded by `defer_s` from the first hold, not from the latest touch - (re-sent every `_TIER_TOUCH_INTERVAL_S`, so it would never run out); - after that `h` is admitted unheld until new evidence arrives. + At most `defer_s` from the first hold. """ defer_s = self.defer_s # `could_serve`: a store-only role indexes stores it never loads. @@ -576,6 +573,24 @@ def state_tier_cpu_bytes() -> int: return int(float(os.environ.get("OFFLOAD_STATE_CPU_SIZE", "32")) * (1 << 30)) +def _lmcache_engine_config(config): + """The LMCache config this engine's offload connector builds; raises when + LMCache or its config cannot be read.""" + from atom.kv_transfer.offload.config import build_lmcache_config + + kvcfg = getattr(config, "kv_transfer_config", None) or {} + # Under `multi` the lmcache.* keys (chunk_size, offload_layout) live on + # the offload SUB-connector, not the composite. Passing the raw + # composite here read a zero chunk grid. Unwrap the sub the same way + # `state_tier_capability` does, so the gate's chunk size and the tier + # the capability check builds stay consistent. + if isinstance(kvcfg, dict) and kvcfg.get("kv_connector") == "multi": + sub, _why = _offload_subconfig(kvcfg) + if sub is not None: + kvcfg = sub + return build_lmcache_config(kvcfg) + + def state_tier_chunk_tokens(config) -> int: """The LMCache chunk size in tokens for the joint KV leg, or 0. @@ -594,19 +609,7 @@ def state_tier_chunk_tokens(config) -> int: disable the joint KV load. """ try: - from atom.kv_transfer.offload.config import build_lmcache_config - - kvcfg = getattr(config, "kv_transfer_config", None) or {} - # Under `multi` the lmcache.* keys (chunk_size, offload_layout) live on - # the offload SUB-connector, not the composite. Passing the raw - # composite here read a zero chunk grid. Unwrap the sub the same way - # `state_tier_capability` does, so the gate's chunk size and the tier - # the capability check builds stay consistent. - if isinstance(kvcfg, dict) and kvcfg.get("kv_connector") == "multi": - sub, _why = _offload_subconfig(kvcfg) - if sub is not None: - kvcfg = sub - return int(build_lmcache_config(kvcfg).chunk_size) + return int(_lmcache_engine_config(config).chunk_size) except Exception: # Blind on purpose: this runs at model load, the import reaches a # third-party package that may be absent entirely, and the only @@ -625,15 +628,7 @@ def state_tier_verified_for_s(config) -> float: when it cannot be read. Like `state_tier_chunk_tokens`, call it only when a tier is hosted.""" try: - from atom.kv_transfer.offload.config import build_lmcache_config - - kvcfg = getattr(config, "kv_transfer_config", None) or {} - # The offload sub-connector's keys, as in `state_tier_chunk_tokens`. - if isinstance(kvcfg, dict) and kvcfg.get("kv_connector") == "multi": - sub, _why = _offload_subconfig(kvcfg) - if sub is not None: - kvcfg = sub - pool_gib = float(build_lmcache_config(kvcfg).max_local_cpu_size) + pool_gib = float(_lmcache_engine_config(config).max_local_cpu_size) except Exception: # noqa: BLE001 - `state_tier_chunk_tokens` logged why logger.warning( "state offload: LMCache's CPU pool size is unknown; a confirmed " diff --git a/tests/conftest.py b/tests/conftest.py index 624f90acff..c50d3c29f0 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -294,6 +294,14 @@ def release_graphs(self): return recordings +@pytest.fixture +def offload_clock(monkeypatch): + """`state_offload`'s monotonic clock, advanced by hand.""" + now = [1000.0] + monkeypatch.setattr("atom.model_engine.state_offload.monotonic", lambda: now[0]) + return now + + @pytest.fixture(autouse=True) def mono_runs_on_gfx950(monkeypatch): """Mono's deployment checks see gfx950, the only arch its kernels are built diff --git a/tests/test_block_manager.py b/tests/test_block_manager.py index ea3853f6db..f3688089fa 100644 --- a/tests/test_block_manager.py +++ b/tests/test_block_manager.py @@ -1393,5 +1393,5 @@ def test_an_unreadable_lmcache_config_gets_the_floor(self, monkeypatch, caplog): with caplog.at_level(logging.WARNING, logger="atom"): bm = _bm_with_state_tier(monkeypatch, _OFFLOAD_KVC) assert bm._joint_chunk_tokens == 0 - assert bm.state_offload.verified_for_s == so._TIER_VERIFIED_MIN_S + assert bm.state_offload.verified_for_s == so.tier_verified_for_s(0) assert sum(r.exc_info is not None for r in caplog.records) == 1 diff --git a/tests/test_lmcache_offload_connector.py b/tests/test_lmcache_offload_connector.py index 0b43aa3eda..9b659297e4 100644 --- a/tests/test_lmcache_offload_connector.py +++ b/tests/test_lmcache_offload_connector.py @@ -46,6 +46,7 @@ SaveOperationId, StateStoreOperationId, StateTouchId, + connector_metadata_has_work, ) from atom.kv_transfer.offload import _block_gpu_connector from atom.kv_transfer.offload import config as offcfg @@ -81,7 +82,10 @@ KimiK3OffloadScheduler, save_stall_seconds, ) -from atom.kv_transfer.offload.hybrid.kimi_k3.state_object import TouchResult +from atom.kv_transfer.offload.hybrid.kimi_k3.state_object import ( + StateByteCodec, + TouchResult, +) from atom.kv_transfer.offload.hybrid.kimi_k3.state_tier import _JointPark from atom.kv_transfer.offload.hybrid.qwen.connector import QwenOffloadScheduler from atom.kv_transfer.offload.metadata import ( @@ -6356,6 +6360,12 @@ def test_a_step_whose_only_work_is_a_state_store_reaches_the_worker(): assert connector_metadata_has_work(meta) +def test_a_step_whose_only_work_is_a_state_touch_reaches_the_worker(): + meta = LMCacheOffloadMetadata() + meta.state_touches = [StateTouchId(7, 0)] + assert connector_metadata_has_work(meta) + + # ── the delegating shells must expose everything their driver reads ──────── # # There are two forwarding axes, not one, and every round of review has found @@ -7264,8 +7274,6 @@ def evict_oldest(self): def _lru_codec(storage, max_entries): - from atom.kv_transfer.offload.hybrid.kimi_k3.state_object import StateByteCodec - codec = StateByteCodec( backend=SimpleNamespace( page_unit_views=lambda units: [], state_entry_views=lambda slot: [] @@ -7308,11 +7316,12 @@ def test_a_touch_moves_the_image_past_older_kv(self, lmcache_key): def test_a_touch_reports_an_evicted_image_gone(self, lmcache_key): pool = _LocalCPU() - codec, _ = _lru_codec(pool, max_entries=2) - assert codec.put(1, []) - assert pool.evict_oldest() == codec.key(1) - assert codec.touch(1) is TouchResult.GONE - assert codec.key(1) not in codec._cache_keys, "no longer counts against the cap" + codec, _ = _lru_codec(pool, max_entries=3) + assert codec.put(1, []) and codec.put(2, []) and codec.put(3, []) + assert pool.remove(codec.key(2)) + assert codec.touch(2) is TouchResult.GONE + assert codec.put(4, []) + assert codec.key(1) in pool.hot, "a gone image no longer counts against the cap" def test_a_resume_refreshes_the_image(self, lmcache_key): """`get_blocking` takes a reference and nothing else.""" @@ -7354,7 +7363,7 @@ def test_the_cap_retires_the_least_recently_used_image(self, lmcache_key): assert codec.key(2) not in pool.hot assert objs[1].refs == 0, "a retired image goes back to the allocator" - assert list(codec._cache_keys) == [codec.key(1), codec.key(3)] + assert codec.key(1) in pool.hot and codec.key(3) in pool.hot @pytest.mark.parametrize( "missing", @@ -7436,8 +7445,21 @@ def _touch_worker(tier): return c +def _reported(worker): + out = KVConnectorOutput() + worker._emit_touch_verdicts(out) + return out.connector_completions + + +def _reports(touch_id, *, resident, refreshed): + return { + ConnectorCompletion(STATE_TOUCH_CHANNEL, touch_id, resident), + ConnectorCompletion(STATE_TOUCH_REFRESH_CHANNEL, touch_id, refreshed), + } + + class TestConnectorTouches: - def test_the_worker_touches_and_records_each_verdict(self): + def test_the_worker_touches_and_reports_each_verdict(self): touched = [] found = {1: TouchResult.REFRESHED, 2: TouchResult.GONE, 3: TouchResult.PRESENT} @@ -7450,18 +7472,34 @@ def touch(h): ids = [StateTouchId(1, 0), StateTouchId(2, 1), StateTouchId(3, 2)] worker._touch_states(SimpleNamespace(state_touches=ids)) assert touched == [1, 2, 3] - # (resident, refreshed): only a present-but-unrefreshed image is False - # on the second. - assert worker._touch_verdicts == { - ids[0]: (True, True), - ids[1]: (False, True), - ids[2]: (True, False), - } + # Only a present-but-unrefreshed image reports unrefreshed. + assert _reported(worker) == ( + _reports(ids[0], resident=True, refreshed=True) + | _reports(ids[1], resident=False, refreshed=True) + | _reports(ids[2], resident=True, refreshed=False) + ) + assert _reported(worker) == set(), "drained once" + + def test_start_load_kv_touches(self, monkeypatch): + monkeypatch.setattr(DenseOffloadConnector, "start_load_kv", lambda s, m: None) + worker = _touch_worker( + SimpleNamespace( + codec=SimpleNamespace(touch=lambda h: TouchResult.REFRESHED) + ) + ) + meta = LMCacheOffloadMetadata() + meta.state_touches = [StateTouchId(7, 0)] + worker.start_load_kv(meta) + assert _reported(worker) == _reports( + StateTouchId(7, 0), resident=True, refreshed=True + ) def test_with_no_tier_every_image_is_gone(self): worker = _touch_worker(None) worker._touch_states(SimpleNamespace(state_touches=[StateTouchId(1, 0)])) - assert worker._touch_verdicts == {StateTouchId(1, 0): (False, True)} + assert _reported(worker) == _reports( + StateTouchId(1, 0), resident=False, refreshed=True + ) def test_a_failing_touch_is_unrefreshed_and_logged_once(self, caplog): """A raising touch keeps the image (resident) but is not a refresh; @@ -7475,7 +7513,9 @@ def touch(h): with caplog.at_level(logging.WARNING, logger="atom"): worker._touch_states(SimpleNamespace(state_touches=ids[:2])) worker._touch_states(SimpleNamespace(state_touches=ids[2:])) - assert worker._touch_verdicts == {i: (True, False) for i in ids} + assert _reported(worker) == set().union( + *(_reports(i, resident=True, refreshed=False) for i in ids) + ) assert [r.exc_info is not None for r in caplog.records] == [True] def _scheduler(self, monkeypatch): @@ -7519,17 +7559,7 @@ def test_both_touch_reports_reach_quorum_across_tp(self, monkeypatch): SimpleNamespace(codec=SimpleNamespace(touch=lambda h, r=result: r)) ) worker._touch_states(SimpleNamespace(state_touches=[touch])) - resident, refreshed = worker._touch_verdicts[touch] - outputs.append( - KVConnectorOutput( - connector_completions={ - ConnectorCompletion(STATE_TOUCH_CHANNEL, touch, resident), - ConnectorCompletion( - STATE_TOUCH_REFRESH_CHANNEL, touch, refreshed - ), - } - ) - ) + outputs.append(KVConnectorOutput(connector_completions=_reported(worker))) merged = KVOutputAggregator(world_size=2).aggregate(outputs) s = self._scheduler(monkeypatch) monkeypatch.setattr( diff --git a/tests/test_offload_shell.py b/tests/test_offload_shell.py index 7d43726817..aa316b9d51 100644 --- a/tests/test_offload_shell.py +++ b/tests/test_offload_shell.py @@ -7,7 +7,13 @@ from types import SimpleNamespace +import pytest + from atom.kv_transfer.disaggregation.factory import KVConnectorFactory +from atom.kv_transfer.disaggregation.multi.multi_connector import ( + MultiConnectorScheduler, +) +from atom.kv_transfer.disaggregation.types import ConnectorCompletion, StateTouchId from atom.kv_transfer.offload import config as offcfg from atom.kv_transfer.offload.connector import ( LMCacheOffloadConnector, @@ -22,6 +28,11 @@ DSV4OffloadConnector, DSV4OffloadScheduler, ) +from atom.kv_transfer.offload.hybrid.kimi_k3.connector import ( + STATE_TOUCH_CHANNEL, + KimiK3OffloadScheduler, +) +from atom.model_engine.state_offload import StateOffloadIndex def _config(*, compress_ratios=None, layout=None): @@ -133,3 +144,32 @@ def test_scheduler_shell_forwards_statistics(): shell._impl = SimpleNamespace(get_statistics=lambda: expected) assert shell.get_statistics() is expected + + +@pytest.mark.parametrize("composite", [False, True], ids=["shell", "multi"]) +def test_state_touches_round_trip_through_the_scheduler_shells( + composite, offload_clock +): + impl = KimiK3OffloadScheduler.__new__(KimiK3OffloadScheduler) + impl._pending_state_touches = [] + impl._touch_generation = 0 + impl._state_touch_verdicts = {} + connector = LMCacheOffloadConnectorScheduler.__new__( + LMCacheOffloadConnectorScheduler + ) + connector._impl = impl + if composite: + shell = connector + connector = MultiConnectorScheduler.__new__(MultiConnectorScheduler) + connector._connectors = [SimpleNamespace(), shell] + idx = StateOffloadIndex(verified_for_s=5.0) + idx.note_stored(7) + offload_clock[0] += 1 + idx.want_touch(7) + idx.publish_touches(connector) + assert idx.refresh_hits, "the touch reached the tier" + impl.connector_completion( + ConnectorCompletion(STATE_TOUCH_CHANNEL, StateTouchId(7, 0), False) + ) + idx.collect_touch_verdicts(connector) + assert not idx.could_serve(7), "the verdict came back" diff --git a/tests/test_state_checkpoint.py b/tests/test_state_checkpoint.py index 672b62d0f5..8bc745f81f 100644 --- a/tests/test_state_checkpoint.py +++ b/tests/test_state_checkpoint.py @@ -18,7 +18,9 @@ import pytest from conftest import MockConfig +from test_state_offload_index import STALE +from atom.kv_transfer.disaggregation.types import KVConnectorOutput from atom.model_engine.block_manager import ( AdmissionProbe, BlockManager, @@ -34,7 +36,7 @@ from atom.model_engine.scheduler import ScheduledBatchOutput, Scheduler from atom.model_engine.sequence import Sequence, SequenceType from atom.model_engine.state_cache import StateCache -from atom.model_engine.state_offload import _TIER_VERIFIED_MAX_S, StateOffloadIndex +from atom.model_engine.state_offload import StateOffloadIndex from atom.model_engine.state_pool import StateSlotPool from atom.model_engine.state_runtime import ( StateRuntime, @@ -4566,18 +4568,6 @@ def test_a_cold_start_without_a_boundary_is_untouched(self): # ── state tier hit refresh: the admission hold and the voters ────────────── -#: Past every confirmation window `tier_verified_for_s` can give. -STALE = _TIER_VERIFIED_MAX_S + 1 - - -@pytest.fixture -def offload_clock(monkeypatch): - """`state_offload`'s monotonic clock, advanced by hand.""" - now = [1000.0] - monkeypatch.setattr("atom.model_engine.state_offload.monotonic", lambda: now[0]) - return now - - def _paged_bm(): bm = make_block_manager( ckpt_config(num_kvcache_blocks=160, pool_entries={"state": 8}), @@ -4844,3 +4834,19 @@ def test_the_scheduler_publishes_touches_with_the_state_loads(offload_clock): s.block_manager = SimpleNamespace(state_offload=idx) s._publish_state_loads() assert sent == [7] + + +def test_the_scheduler_applies_touch_verdicts(offload_clock): + idx = StateOffloadIndex(verified_for_s=5.0) + idx.note_stored(7) + offload_clock[0] += 1 + idx.want_touch(7) + assert idx.take_touches() == [7] + s = Scheduler.__new__(Scheduler) + s.kv_connector = SimpleNamespace(take_state_touch_verdicts=lambda: {7: False}) + s.block_manager = SimpleNamespace(state_offload=idx) + s.finished_recving_kv_req_ids = [] + s.failed_recving_kv_req_ids = [] + s.deferred_free_blocks = {} + s._update_from_kv_xfer_finished(KVConnectorOutput()) + assert not idx.could_serve(7) diff --git a/tests/test_state_offload_index.py b/tests/test_state_offload_index.py index e9be6c39bf..aa70c86376 100644 --- a/tests/test_state_offload_index.py +++ b/tests/test_state_offload_index.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: MIT """StateOffloadIndex must stop advertising images the worker codec dropped.""" +from math import inf from types import SimpleNamespace import pytest @@ -60,18 +61,10 @@ def test_state_tier_cpu_bytes_reads_gib(monkeypatch): #: Past every confirmation window `tier_verified_for_s` can give. -STALE = so._TIER_VERIFIED_MAX_S + 1 +STALE = so.tier_verified_for_s(inf) + 1 -@pytest.fixture -def offload_clock(monkeypatch): - """`state_offload`'s monotonic clock, advanced by hand.""" - now = [1000.0] - monkeypatch.setattr("atom.model_engine.state_offload.monotonic", lambda: now[0]) - return now - - -def _index(*hashes, verified_for_s=so._TIER_VERIFIED_MIN_S, **kw): +def _index(*hashes, verified_for_s=5.0, **kw): idx = StateOffloadIndex(verified_for_s=verified_for_s, **kw) for h in hashes: idx.note_stored(h) @@ -94,7 +87,10 @@ def test_no_touches_without_refresh(self, offload_clock): idx = _index(7, verified_for_s=None) idx.want_touch(7) assert idx.take_touches() == [] - assert idx._verified_at == {}, "nothing reads it" + offload_clock[0] += STALE + assert not idx.hold_admission(7) + # Up to the hash cap of entries nothing would read. + assert idx._verified_at == {} def test_a_hash_voted_every_pass_is_touched_at_an_interval(self, offload_clock): idx = _index(7) @@ -278,15 +274,6 @@ class TestTouchPlumbing: """The scheduler's two calls a pass: `publish_touches` before the metadata is built, `collect_touch_verdicts` before the next admission.""" - def test_voted_hashes_reach_the_connector(self, offload_clock): - idx = _index(7) - idx.want_touch(7) - sent = [] - idx.publish_touches( - SimpleNamespace(enqueue_state_touches=lambda hs: sent.extend(hs) or True) - ) - assert sent == [7] and idx.refresh_hits - def test_a_connector_without_touches_turns_the_refresh_off(self, offload_clock): idx = _index(7) idx.want_touch(7) @@ -295,14 +282,3 @@ def test_a_connector_without_touches_turns_the_refresh_off(self, offload_clock): assert idx.take_touches() == [] offload_clock[0] += STALE assert not idx.hold_admission(7), "nothing would ever answer the hold" - - def test_verdicts_are_collected(self, offload_clock): - idx = _index(7, 8) - offload_clock[0] += 1 - idx.want_touch(7) - idx.want_touch(8) - assert sorted(idx.take_touches()) == [7, 8] - idx.collect_touch_verdicts( - SimpleNamespace(take_state_touch_verdicts=lambda: {7: False, 8: True}) - ) - assert not idx.could_serve(7) and idx.could_serve(8)