From 5e4f77e86250bae7aa4a04cf596403b3a950a95d Mon Sep 17 00:00:00 2001
From: RunningLeon <mnsheng@yeah.net>
Date: Fri, 16 Oct 2026 09:04:13 +0000
Subject: [PATCH 1/3] fix(pytorch): invalidate local prefix cache on sleep

---
 lmdeploy/pytorch/engine/engine.py             |  2 +-
 lmdeploy/pytorch/paging/block_trie/trie.py    | 25 +++++++
 lmdeploy/pytorch/paging/scheduler.py          | 11 +++
 tests/pytorch/engine/test_engine_sleep.py     | 70 +++++++++++++++--
 .../paging/test_block_trie/test_clear.py      | 75 +++++++++++++++++++
 .../paging/test_scheduler_kv_transfer.py      | 45 ++++++++++-
 6 files changed, 220 insertions(+), 8 deletions(-)
 create mode 100644 tests/pytorch/paging/test_block_trie/test_clear.py

diff --git a/lmdeploy/pytorch/engine/engine.py b/lmdeploy/pytorch/engine/engine.py
index 22066bebb3..1fc30dc2b8 100644
--- a/lmdeploy/pytorch/engine/engine.py
+++ b/lmdeploy/pytorch/engine/engine.py
@@ -623,7 +623,7 @@ async def sleep(self, level: int = 1):
         # cancel all remain sessions
         self._cancel_and_end_all_sessions()
         await self.executor.sleep(level)
-        self.scheduler.finish_kv_transfers_after_worker_drain()
+        self.scheduler.reset_cache_after_worker_drain()
         if self._engine_loop is not None:
             self._engine_loop.reset_runtime_state()
         logger.info('PyTorch engine entered sleep: level=%s, sleeping_tags=%s.', level, sorted(self._sleeping_tags))
diff --git a/lmdeploy/pytorch/paging/block_trie/trie.py b/lmdeploy/pytorch/paging/block_trie/trie.py
index 423ad9fcc3..6516b0fc41 100644
--- a/lmdeploy/pytorch/paging/block_trie/trie.py
+++ b/lmdeploy/pytorch/paging/block_trie/trie.py
@@ -673,6 +673,31 @@ def allocate(self, seq: SchedulerSequence):
         self._extend_trie_path(seq, node, num_complete_blocks)
         recompute_overlap.clear_fresh_block_range()
 
+    def clear(self) -> None:
+        """Release all local prefix ownership after requests and workers drain.
+
+        Unlike pressure eviction, this also releases root-backed partial
+        checkpoints and unpublished reservations. Refuse live references/pins
+        before mutating anything: clearing identity cannot make in-flight
+        cache copies safe.
+        """
+        nodes = list(self._roots.values())
+        for node in nodes:
+            nodes.extend(node.children.values())
+        if any(self._state_checkpoints.is_pinned(node) for node in nodes):
+            raise RuntimeError('Cannot clear pinned prefix-cache checkpoints.')
+        blocks = np.array([node.block_id for node in nodes if node.block_id >= 0], dtype=np.int64)
+        if len(blocks) and np.any(self.allocator.get_ref_count(blocks) != 1):
+            raise RuntimeError('Cannot clear referenced prefix-cache blocks.')
+
+        for node in reversed(nodes):
+            self._state_checkpoints.release_checkpoint(node)
+            node.detach_leaf()
+        if len(blocks):
+            self.allocator.free(blocks)
+        self._roots.clear()
+        self._kv_lifecycle.leaves.clear()
+
     def evict(self, max_num_blocks: int):
         """Evict checkpoint-frozen and trie-owned KV blocks.
 
diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index 714228094d..135d375615 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -501,6 +501,17 @@ def finish_kv_transfers_after_worker_drain(self) -> None:
             self.kv_connector.finish_transfers_after_worker_drain()
         self.kv_save_coordinator.clear()
 
+    def reset_cache_after_worker_drain(self) -> None:
+        """Invalidate local cache identity after sleep discards worker storage.
+
+        Ended external loads and save leases must release their references first. Remote Store objects are not
+        invalidated; future requests may restore them through the normal lookup/load completion protocol.
+        """
+        self.finish_kv_transfers_after_worker_drain()
+        if self.sessions:
+            raise RuntimeError('Cannot reset cache before all sessions have ended.')
+        self.block_trie.clear()
+
     def get_block_tables(self, seqs: SeqList):
         """Get block tables for the sequences."""
         return [self.block_manager.get_block_table(seq) for seq in seqs]
diff --git a/tests/pytorch/engine/test_engine_sleep.py b/tests/pytorch/engine/test_engine_sleep.py
index ae2ffd7245..5223b34dc8 100644
--- a/tests/pytorch/engine/test_engine_sleep.py
+++ b/tests/pytorch/engine/test_engine_sleep.py
@@ -3,12 +3,17 @@
 from types import SimpleNamespace
 
 import pytest
+import torch
 
 from lmdeploy.messages import EngineOutput, ResponseType
+from lmdeploy.pytorch.config import CacheConfig, SchedulerConfig
 from lmdeploy.pytorch.engine.engine import Engine
 from lmdeploy.pytorch.engine.engine_instance import EngineInstance
 from lmdeploy.pytorch.engine.executor.mp_executor import MPExecutor
 from lmdeploy.pytorch.engine.request import RequestManager, RequestType, Response
+from lmdeploy.pytorch.messages import SequenceMeta
+from lmdeploy.pytorch.paging import Scheduler
+from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
 
 
 class _FakeSequence:
@@ -26,9 +31,10 @@ def __init__(self, session_id, seq):
 
 class _FakeScheduler:
 
-    def __init__(self, session):
+    def __init__(self, session, events):
         self.sessions = {1: session}
         self.ended_sessions = []
+        self.events = events
 
     def end_session(self, session_id):
         self.ended_sessions.append(session_id)
@@ -40,8 +46,9 @@ def get_session(self, session_id):
     def get_sessions(self):
         return list(self.sessions.values())
 
-    def finish_kv_transfers_after_worker_drain(self):
-        pass
+    def reset_cache_after_worker_drain(self):
+        assert self.events[-1] == 'sleep'
+        self.events.append('reset_cache')
 
 
 class _FakeEngineLoop:
@@ -112,9 +119,9 @@ def _build_sleeping_test_engine(event_loop):
     resp = Response(type=ResponseType.INTERNAL_ENGINE_ERROR, sender_id=0, event=asyncio.Event())
     seq = _FakeSequence(resp)
     session = _FakeSession(1, seq)
-    engine.scheduler = _FakeScheduler(session)
-    engine._sleeping_tags = set()
     engine.events = []
+    engine.scheduler = _FakeScheduler(session, engine.events)
+    engine._sleeping_tags = set()
     engine._engine_loop = _FakeEngineLoop(engine)
     engine.executor = _FakeExecutor(engine)
     return engine, resp
@@ -133,7 +140,7 @@ def test_engine_sleep_blocks_inputs_cancels_sessions_then_sleeps(event_loop):
     assert resp.event.is_set()
     assert engine.scheduler.ended_sessions == [1]
     assert engine.executor.sleep_calls == [1]
-    assert engine.events == ['drain', 'sleep', 'reset_engine_loop']
+    assert engine.events == ['drain', 'sleep', 'reset_cache', 'reset_engine_loop']
 
 
 def test_engine_wakeup_reenables_inputs_only_after_all_tags(event_loop):
@@ -175,6 +182,57 @@ def test_engine_wakeup_all_delegates_warmup_to_executor_wakeup(event_loop):
     assert engine.events == [('wakeup', None), 'resume']
 
 
+@pytest.mark.parametrize('level', [1, 2])
+@pytest.mark.parametrize('staged', [False, True])
+@pytest.mark.parametrize('state_step', [None, 8, 32, 35])
+@pytest.mark.parametrize('prefix_caching', [False, True])
+def test_sleep_invalidates_completed_prefix_cache(event_loop, level, staged, state_step, prefix_caching):
+    cache = CacheConfig(max_batches=4,
+                        block_size=16,
+                        num_cpu_blocks=4,
+                        num_gpu_blocks=16,
+                        enable_prefix_caching=prefix_caching,
+                        num_state_caches=13,
+                        prefix_cache_state_budget=8,
+                        states_shapes=[] if state_step is None else [((1, ), torch.float32)])
+    scheduler = Scheduler(SchedulerConfig(max_batches=4, max_session_len=128, max_request_output_len=64),
+                          cache, seq_meta=SequenceMeta(16, strategy=ARSequenceStrategy()))
+    engine, _ = _build_sleeping_test_engine(event_loop)
+    engine.scheduler = scheduler
+    step = 32 if state_step is None else state_step
+    tokens = [1] * step
+    for _ in range(3):
+        seq = scheduler.add_session(0).add_sequence(tokens)
+        scheduler.block_manager.allocate(seq)
+        scheduler.block_trie.allocate(seq)
+        if state_step is not None and prefix_caching:
+            assert scheduler.state_checkpoints.reserve_save(seq) >= 0
+            assert scheduler.state_checkpoints.publish_save(seq)
+        scheduler.end_session(0)
+        repeat = scheduler.add_session(1).add_sequence(tokens + [2])
+        scheduler.block_trie.match(repeat)
+        assert repeat.num_history_ids == (step if prefix_caching else 0)
+        # Leave a matched request alive, including its restore pin, for sleep
+        # to cancel in addition to the already completed producer.
+        scheduler.state_checkpoints.pin_restore(repeat)
+        event_loop.run_until_complete(engine.sleep(level))
+        assert scheduler.block_manager.get_num_free_gpu_blocks() == cache.num_gpu_blocks
+        assert scheduler.state_manager.get_num_free() == cache.num_state_caches - 1
+        if staged:
+            engine.wakeup(['weights'])
+            assert engine.req_manager.is_request_blocked(RequestType.ADD_MESSAGE)
+            engine.wakeup(['kv_cache'])
+        else:
+            engine.wakeup()
+        assert not engine.req_manager.is_request_blocked(RequestType.ADD_MESSAGE)
+        new = scheduler.add_session(2).add_sequence(tokens + [2])
+        scheduler.block_trie.match(new)
+        assert new.num_history_ids == 0
+        assert new.cached_tokens == 0
+        assert not new.prefix_cache.restore.is_selected
+        scheduler.end_session(2)
+
+
 def test_complete_weights_update_keeps_kv_cache_sleeping(event_loop):
     engine, _ = _build_sleeping_test_engine(event_loop)
     engine._sleeping_tags = {'weights', 'kv_cache'}
diff --git a/tests/pytorch/paging/test_block_trie/test_clear.py b/tests/pytorch/paging/test_block_trie/test_clear.py
new file mode 100644
index 0000000000..3d1e9b1f51
--- /dev/null
+++ b/tests/pytorch/paging/test_block_trie/test_clear.py
@@ -0,0 +1,75 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import pytest
+
+
+@pytest.mark.parametrize('step', [8, 32, 35])
+@pytest.mark.parametrize('published', [False, True])
+def test_clear_releases_checkpoints_and_all_kv(ssm_scheduler, step, published):
+    scheduler = ssm_scheduler
+    trie = scheduler.block_trie
+    manager = scheduler.block_manager
+    initial_free = manager.get_num_free_gpu_blocks()
+    initial_logical_free = manager.allocator.get_num_free_blocks()
+    initial_states = scheduler.state_manager.get_num_free()
+    nodes = []
+    for session_id, adapter in enumerate([None, 'adapter']):
+        seq = scheduler.add_session(session_id).add_sequence([1] * step, adapter_name=adapter)
+        manager.allocate(seq)
+        trie.allocate(seq)
+        assert trie.state_checkpoints.reserve_save(seq) >= 0
+        nodes.append(seq.prefix_cache.pending_save.node)
+        if published:
+            assert trie.state_checkpoints.publish_save(seq)
+        # Keep unpublished reservations to exercise clear's ownership rather
+        # than the usual end-session discard path.
+        seq.prefix_cache.pending_save.clear()
+        scheduler.end_session(session_id)
+
+    trie.clear()
+    trie.clear()  # Idempotent; no duplicate frees.
+
+    assert manager.get_num_free_gpu_blocks() == initial_free
+    assert manager.allocator.get_num_free_blocks() == initial_logical_free
+    assert scheduler.state_manager.get_num_free() == initial_states
+    assert not trie.leaves
+    assert not list(trie._checkpoint_index.unique_nodes())
+    assert all(node.state_checkpoint is None for node in nodes)
+    for session_id, adapter in enumerate([None, 'adapter']):
+        seq = scheduler.add_session(session_id).add_sequence([1] * step + [2], adapter_name=adapter)
+        trie.match(seq)
+        assert seq.num_history_ids == 0
+        assert not seq.prefix_cache.restore.is_selected
+
+
+def test_clear_rejects_live_kv_refs_without_partial_release(scheduler):
+    trie = scheduler.block_trie
+    seq = scheduler.add_session(0).add_sequence([1] * 32)
+    scheduler.block_manager.allocate(seq)
+    trie.allocate(seq)
+    free = scheduler.block_manager.get_num_free_gpu_blocks()
+
+    with pytest.raises(RuntimeError, match='referenced'):
+        trie.clear()
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == free
+    scheduler.end_session(0)
+    trie.clear()
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == scheduler.cache_config.num_gpu_blocks
+
+
+def test_clear_rejects_checkpoint_copy_pins(ssm_scheduler):
+    scheduler = ssm_scheduler
+    trie = scheduler.block_trie
+    seq = scheduler.add_session(0).add_sequence([1] * 32)
+    scheduler.block_manager.allocate(seq)
+    trie.allocate(seq)
+    trie.state_checkpoints.reserve_save(seq)
+    trie.state_checkpoints.publish_save(seq, pin_save=True)
+    scheduler.end_session(0)
+    free = scheduler.block_manager.get_num_free_gpu_blocks()
+
+    with pytest.raises(RuntimeError, match='pinned'):
+        trie.clear()
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == free
+    trie.state_checkpoints.unpin_save(seq)
+    trie.clear()
+    assert scheduler.state_manager.get_num_allocated_checkpoint_states() == 0
diff --git a/tests/pytorch/paging/test_scheduler_kv_transfer.py b/tests/pytorch/paging/test_scheduler_kv_transfer.py
index 45e563a546..9a15dec7c0 100644
--- a/tests/pytorch/paging/test_scheduler_kv_transfer.py
+++ b/tests/pytorch/paging/test_scheduler_kv_transfer.py
@@ -793,11 +793,12 @@ def test_worker_drain_finishes_an_ended_session_with_a_dropped_load_output():
     scheduler.schedule(is_prefill=True)
     scheduler.end_session(85)
 
-    scheduler.finish_kv_transfers_after_worker_drain()
+    scheduler.reset_cache_after_worker_drain()
 
     assert 85 not in scheduler.sessions
     assert connector.finished == [seq.seq_id]
     assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 0
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == scheduler.cache_config.num_gpu_blocks
 
 
 def test_completed_async_load_is_admitted_before_an_older_waiter():
@@ -928,3 +929,45 @@ def get_save_block_leases(self):
 
     assert scheduler.block_manager.get_num_free_gpu_blocks() == 2
     assert not scheduler.kv_save_coordinator.has_pending()
+
+
+def test_cache_reset_drains_save_leases_before_clearing_local_prefix():
+
+    class _SaveMetadata(KVConnectorMetadata):
+
+        def get_save_block_leases(self):
+            return (KVSaveBlockLease(9, tuple(logical_blocks)), )
+
+    connector = _AsyncLookupConnector([])
+    scheduler = _make_async_lookup_scheduler(connector, num_gpu_blocks=4)
+    seq = scheduler.add_session(90).add_sequence(torch.arange(8))
+    scheduler.block_manager.allocate(seq)
+    scheduler.block_trie.allocate(seq)
+    logical_blocks = seq.logical_blocks.get_real_blocks().copy()
+    connector.build_connector_meta = lambda step_input: _SaveMetadata()
+    scheduler.build_connector_meta([seq], connector_token_lens=(8, ))
+    scheduler.end_session(90)
+    assert scheduler.kv_save_coordinator.has_pending()
+    assert all(scheduler.block_manager.allocator.get_ref_count(logical_blocks) == 2)
+
+    def finish_transfers():
+        # Both the trie and asynchronous writer still own the blocks until
+        # worker drain is acknowledged. Do not delete remote Store contents.
+        assert all(scheduler.block_manager.allocator.get_ref_count(logical_blocks) == 2)
+
+    connector.finish_transfers_after_worker_drain = finish_transfers
+    scheduler.reset_cache_after_worker_drain()
+
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == 4
+    assert not scheduler.kv_save_coordinator.has_pending()
+    assert connector.shutdown_calls == 0
+    newcomer = scheduler.add_session(91).add_sequence(torch.arange(9))
+    scheduler.block_trie.match(newcomer)
+    assert newcomer.num_history_ids == 0
+
+
+def test_cache_reset_rejects_live_sessions():
+    scheduler = _make_async_lookup_scheduler(_AsyncLookupConnector([]))
+    scheduler.add_session(92)
+    with pytest.raises(RuntimeError, match='sessions'):
+        scheduler.reset_cache_after_worker_drain()

From 3b189d07ee505ba280cb24910a7f89ff423ff565 Mon Sep 17 00:00:00 2001
From: RunningLeon <mnsheng@yeah.net>
Date: Fri, 16 Oct 2026 10:48:44 +0000
Subject: [PATCH 2/3] fix(pytorch): invalidate remote Store cache on sleep

---
 lmdeploy/pytorch/engine/model_agent/agent.py  | 19 +++++++
 lmdeploy/pytorch/kv_connector/factory.py      |  2 +
 .../kv_connector/mooncake/store/connector.py  |  2 +
 .../kv_connector/mooncake/store/data.py       |  6 ++-
 .../kv_connector/mooncake/store/worker.py     |  2 +
 lmdeploy/pytorch/paging/scheduler.py          |  4 +-
 .../engine/test_model_agent_kv_connector.py   | 49 +++++++++++++++++--
 .../test_mooncake_store_connector.py          |  2 +
 .../test_mooncake_store_worker.py             | 29 ++++++++++-
 9 files changed, 108 insertions(+), 7 deletions(-)

diff --git a/lmdeploy/pytorch/engine/model_agent/agent.py b/lmdeploy/pytorch/engine/model_agent/agent.py
index 511efac8e4..9e743a1cf3 100644
--- a/lmdeploy/pytorch/engine/model_agent/agent.py
+++ b/lmdeploy/pytorch/engine/model_agent/agent.py
@@ -10,6 +10,7 @@
 from multiprocessing.reduction import ForkingPickler
 from os import getenv
 from typing import Any
+from uuid import uuid4
 
 import numpy as np
 import pybase64
@@ -384,6 +385,7 @@ def __init__(
         self._cache_plan_block_nbytes: tuple[int, int]
         self.state_cache_engine = None
         self.kv_connector = None
+        self._kv_cache_generation = ''
         self.profiler: AgentProfiler = None
         try:
             self.guided_decoding_manager = GuidedDecodingManager(self.tokenizer, model_config.vocab_size)
@@ -1465,6 +1467,19 @@ def shutdown_kv_connector(self):
         self._release_completed_h2d_transfers()
         self._shutdown_kv_connector()
 
+    def _advance_kv_cache_generation(self):
+        """Isolate post-sleep remote KV from any previous weights or writes.
+
+        All DP/TP workers use the same fresh namespace. Keep it on the agent so rebuilding models or connectors during
+        staged wakeup cannot lose it. Old Store objects are left for eviction, not deleted from shared storage.
+        """
+        generation = [uuid4().hex if self.rank == 0 else None]
+        if self.dist_config.world_size > 1:
+            dist.broadcast_object_list(generation, src=0, group=self.dist_ctx.cpu_group)
+        self._kv_cache_generation = generation[0]
+        logger.info('Reset external KV cache generation: rank=%d generation=%s',
+                    self.rank, self._kv_cache_generation)
+
     def _shutdown_kv_connector(self):
         """Shutdown the connector before releasing its registered caches."""
         if self.kv_connector is None:
@@ -1542,6 +1557,7 @@ def build_cache_engine(self):
                 tp_rank=tp_rank,
                 tp_size=tp,
                 kv_head_replica_num=self.model_config.num_replicate_key_value_heads,
+                cache_generation=self._kv_cache_generation,
             )
             if self.kv_connector is not None:
                 self.kv_connector.register_kv_caches(self._get_connector_kv_caches())
@@ -1913,7 +1929,10 @@ async def sleep(self, level: int = 1):
         device = 'cpu' if level == 1 else 'meta'
         # Stop producers of GPU work and wait for queued work before closing the
         # connector threads that still reference registered cache addresses.
+        had_kv_connector = self.kv_connector is not None
         self.shutdown_kv_connector()
+        if had_kv_connector:
+            self._advance_kv_cache_generation()
         self.cache_engine = None
         self.state_cache_engine = None
         self.reset_graph_runner()
diff --git a/lmdeploy/pytorch/kv_connector/factory.py b/lmdeploy/pytorch/kv_connector/factory.py
index 8c9da3b1d5..c93b33fcdc 100644
--- a/lmdeploy/pytorch/kv_connector/factory.py
+++ b/lmdeploy/pytorch/kv_connector/factory.py
@@ -74,6 +74,7 @@ def build_kv_connector(
     tp_rank: int = 0,
     tp_size: int = 1,
     kv_head_replica_num: int = 1,
+    cache_generation: str = '',
 ) -> KVConnectorBase | None:
     """Build the configured connector without importing optional backends
     eagerly."""
@@ -93,6 +94,7 @@ def build_kv_connector(
             tp_rank=tp_rank,
             tp_size=tp_size,
             kv_head_replica_num=kv_head_replica_num,
+            cache_generation=cache_generation,
         )
 
     raise ValueError(f'Unsupported KV connector: {connector_name!r}')
diff --git a/lmdeploy/pytorch/kv_connector/mooncake/store/connector.py b/lmdeploy/pytorch/kv_connector/mooncake/store/connector.py
index 786c633ae0..a00765e835 100644
--- a/lmdeploy/pytorch/kv_connector/mooncake/store/connector.py
+++ b/lmdeploy/pytorch/kv_connector/mooncake/store/connector.py
@@ -37,6 +37,7 @@ def __init__(
         tp_rank: int = 0,
         tp_size: int = 1,
         kv_head_replica_num: int = 1,
+        cache_generation: str = '',
     ) -> None:
         super().__init__(role)
 
@@ -61,6 +62,7 @@ def __init__(
                 tp_rank=tp_rank,
                 tp_size=tp_size,
                 kv_head_replica_num=kv_head_replica_num,
+                cache_generation=cache_generation,
             )
 
     def _require_scheduler(self) -> MooncakeStoreScheduler:
diff --git a/lmdeploy/pytorch/kv_connector/mooncake/store/data.py b/lmdeploy/pytorch/kv_connector/mooncake/store/data.py
index 8ee9d22b4f..7ca094a128 100644
--- a/lmdeploy/pytorch/kv_connector/mooncake/store/data.py
+++ b/lmdeploy/pytorch/kv_connector/mooncake/store/data.py
@@ -169,6 +169,7 @@ class MooncakeStoreKeyMetadata:
     tp_size: int
     block_size: int
     kv_head_replica_num: int = 1
+    cache_generation: str = ''
 
     def __post_init__(self) -> None:
         validate_kv_head_replica_num(self.kv_head_replica_num, self.tp_size)
@@ -190,8 +191,11 @@ def build_store_key(
         raise ValueError(f'block_hash must contain {MOONCAKE_BLOCK_HASH_BYTES} bytes')
 
     prefix = f'{metadata.cache_prefix}@' if metadata.cache_prefix else ''
+    # Preserve the initial shared namespace, but never reuse pre-sleep KV after
+    # wakeup. Lookup, load and save all construct keys through this function.
+    generation = f'@generation:{metadata.cache_generation}' if metadata.cache_generation else ''
     return (
-        f'{prefix}{metadata.model_name}'
+        f'{prefix}{metadata.model_name}{generation}'
         f'@tp_rank:{kv_head_rank}'
         '@group:0'
         f'@{block_hash.hex()}'
diff --git a/lmdeploy/pytorch/kv_connector/mooncake/store/worker.py b/lmdeploy/pytorch/kv_connector/mooncake/store/worker.py
index f0270c636c..9e7c9d8703 100644
--- a/lmdeploy/pytorch/kv_connector/mooncake/store/worker.py
+++ b/lmdeploy/pytorch/kv_connector/mooncake/store/worker.py
@@ -71,6 +71,7 @@ def __init__(
         tp_rank: int = 0,
         tp_size: int = 1,
         kv_head_replica_num: int = 1,
+        cache_generation: str = '',
         store_factory: StoreFactory | None = None,
         replicate_config: Any = None,
     ) -> None:
@@ -110,6 +111,7 @@ def __init__(
             tp_size=tp_size,
             block_size=cache_config.block_size,
             kv_head_replica_num=kv_head_replica_num,
+            cache_generation=cache_generation,
         )
 
         config_path = extra_config.get('mooncake_config_path')
diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index 135d375615..89158fdfd6 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -504,8 +504,8 @@ def finish_kv_transfers_after_worker_drain(self) -> None:
     def reset_cache_after_worker_drain(self) -> None:
         """Invalidate local cache identity after sleep discards worker storage.
 
-        Ended external loads and save leases must release their references first. Remote Store objects are not
-        invalidated; future requests may restore them through the normal lookup/load completion protocol.
+        Ended external loads and save leases must release their references first. Workers have already advanced the
+        external cache generation, so neither local nor remote pre-sleep KV can be reused after wakeup.
         """
         self.finish_kv_transfers_after_worker_drain()
         if self.sessions:
diff --git a/tests/pytorch/engine/test_model_agent_kv_connector.py b/tests/pytorch/engine/test_model_agent_kv_connector.py
index 6461c571f6..3e5d7fb8c7 100644
--- a/tests/pytorch/engine/test_model_agent_kv_connector.py
+++ b/tests/pytorch/engine/test_model_agent_kv_connector.py
@@ -79,6 +79,7 @@ def fake_connector(role, cache_config, **kwargs):
         'tp_rank': 3,
         'tp_size': 8,
         'kv_head_replica_num': 1,
+        'cache_generation': '',
     }
 
 
@@ -88,6 +89,7 @@ def _bare_model_agent():
     agent = BaseModelAgent.__new__(BaseModelAgent)
     agent.all_context = nullcontext
     agent.cache_config = _enabled_config()
+    agent._kv_cache_generation = 'post-sleep-generation'
     agent.model_config = SimpleNamespace(num_replicate_key_value_heads=4)
     agent.rank = 7
     agent.cache_stream = object()
@@ -171,6 +173,7 @@ def fake_build_connector(role, cache_config, **kwargs):
         'tp_rank': 3,
         'tp_size': 8,
         'kv_head_replica_num': 4,
+        'cache_generation': 'post-sleep-generation',
     }
     assert events[5] == (
         'register', {
@@ -280,7 +283,8 @@ def test_build_cache_engine_rejects_mismatched_mtp_parallel_config(monkeypatch):
     assert agent.kv_connector is None
 
 
-def test_sleep_shuts_down_connector_before_dropping_cache(monkeypatch):
+@pytest.mark.parametrize('level', [1, 2])
+def test_sleep_shuts_down_connector_before_dropping_cache(monkeypatch, level):
     from lmdeploy.pytorch.engine.model_agent import agent as agent_module
 
     events = []
@@ -296,7 +300,8 @@ def shutdown(self):
     agent.kv_connector = _Connector()
     agent.cache_engine = cache_engine
     agent.state_cache_engine = object()
-    agent.dist_config = SimpleNamespace(dp=1)
+    agent.rank = 0
+    agent.dist_config = SimpleNamespace(dp=1, world_size=1)
     agent.state = SimpleNamespace(is_sleeping=False,
                                   to_sleep=SimpleNamespace(clear=lambda: events.append('sleep-clear')))
     model = SimpleNamespace(to=lambda **kwargs: events.append('model-to'))
@@ -312,13 +317,22 @@ def shutdown(self):
     monkeypatch.setattr(agent_module.torch.cuda, 'synchronize', lambda: events.append('cuda-sync'))
     monkeypatch.setattr(agent_module.torch.cuda, 'empty_cache', lambda: events.append('empty-cache'))
 
-    asyncio.run(agent.sleep())
+    def new_generation():
+        assert events[-1] == 'shutdown'
+        assert agent.cache_engine is cache_engine
+        events.append('new-generation')
+        return SimpleNamespace(hex='fresh-generation')
+
+    monkeypatch.setattr(agent_module, 'uuid4', new_generation)
+    asyncio.run(agent.sleep(level))
+    assert agent._kv_cache_generation == 'fresh-generation'
 
     assert events == [
         'drain',
         'cuda-sync',
         'release-h2d',
         'shutdown',
+        'new-generation',
         'reset-graph',
         'model-to',
         'spec-model-to',
@@ -672,3 +686,32 @@ def test_wakeup_rebuilds_connector_with_kv_cache():
 
     assert events == ['build-cache', 'warmup']
     assert agent.state.is_sleeping is False
+
+
+def test_cache_generation_is_shared_across_workers_and_changes_each_cycle(monkeypatch):
+    from lmdeploy.pytorch.engine.model_agent import agent as agent_module
+
+    group = object()
+    shared = []
+    agents = [_bare_model_agent() for _ in range(4)]
+    for rank, agent in enumerate(agents):
+        agent.rank = rank
+        agent.dist_config = SimpleNamespace(world_size=4)
+        agent.dist_ctx = SimpleNamespace(cpu_group=group)
+
+    def broadcast(values, src, group):
+        assert src == 0 and group is agents[0].dist_ctx.cpu_group
+        if values[0] is not None:
+            shared[:] = values
+        else:
+            values[:] = shared
+
+    monkeypatch.setattr(agent_module.dist, 'broadcast_object_list', broadcast)
+    generations = set()
+    for _ in range(3):
+        for agent in agents:
+            agent._advance_kv_cache_generation()
+        generation = agents[0]._kv_cache_generation
+        assert generation and generation not in generations
+        assert all(agent._kv_cache_generation == generation for agent in agents)
+        generations.add(generation)
diff --git a/tests/pytorch/kv_connector/test_mooncake_store_connector.py b/tests/pytorch/kv_connector/test_mooncake_store_connector.py
index 9ab1ba9a5f..7223b7c0d0 100644
--- a/tests/pytorch/kv_connector/test_mooncake_store_connector.py
+++ b/tests/pytorch/kv_connector/test_mooncake_store_connector.py
@@ -73,6 +73,7 @@ def test_connector_constructs_only_worker_delegate(cache_config):
         global_rank=7,
         tp_rank=3,
         tp_size=8,
+        cache_generation='after-sleep',
     )
 
     assert connector.connector_scheduler is None
@@ -82,6 +83,7 @@ def test_connector_constructs_only_worker_delegate(cache_config):
     assert connector.connector_worker.tp_rank == 3
     assert connector.connector_worker.tp_size == 8
     assert connector.kv_role == 'kv_both'
+    assert connector.connector_worker.key_metadata.cache_generation == 'after-sleep'
 
 
 def test_connector_requires_enabled_kv_transfer_config():
diff --git a/tests/pytorch/kv_connector/test_mooncake_store_worker.py b/tests/pytorch/kv_connector/test_mooncake_store_worker.py
index e764ab3a4a..843bc22034 100644
--- a/tests/pytorch/kv_connector/test_mooncake_store_worker.py
+++ b/tests/pytorch/kv_connector/test_mooncake_store_worker.py
@@ -634,7 +634,8 @@ def test_registration_rejects_non_cuda_noncontiguous_and_empty_rows(tmp_path):
     worker.shutdown()
 
 
-def test_async_load_writes_allocated_blocks_and_reports_partial_failure(tmp_path):
+@pytest.mark.parametrize('cache_generation', ['', 'after-sleep'])
+def test_async_load_writes_allocated_blocks_and_reports_partial_failure(tmp_path, cache_generation):
     path = write_store_config(tmp_path)
     store = FakeStore(get_results=[0, -1])
     worker = MooncakeStoreWorker(
@@ -643,6 +644,7 @@ def test_async_load_writes_allocated_blocks_and_reports_partial_failure(tmp_path
         tp_rank=5,
         tp_size=8,
         kv_head_replica_num=4,
+        cache_generation=cache_generation,
         store_factory=lambda: store,
     )
     worker.register_kv_caches({
@@ -724,9 +726,11 @@ def test_async_load_writes_target_and_mtp_cache_rows(tmp_path, monkeypatch):
     worker.shutdown()
 
 
+@pytest.mark.parametrize('cache_generation', ['', 'after-sleep'])
 def test_async_save_waits_for_forward_and_writes_only_owned_missing_blocks(
     tmp_path,
     monkeypatch,
+    cache_generation,
 ):
     path = write_store_config(tmp_path)
     event = FakeCudaEvent()
@@ -746,6 +750,7 @@ def assert_forward_is_ready():
         tp_rank=3,
         tp_size=4,
         kv_head_replica_num=2,
+        cache_generation=cache_generation,
         store_factory=lambda: store,
         replicate_config=replicate_config,
     )
@@ -840,3 +845,25 @@ def test_shutdown_is_idempotent_and_propagates_close_error(tmp_path):
 
     assert store.close_calls == 1
     assert worker.store is None
+
+
+def test_cache_generation_invalidates_remote_hits_without_deleting_shared_objects(tmp_path):
+    # Store storage outlives workers. Even delayed writes to an old generation
+    # cannot be mistaken for KV produced after a sleep/weight update.
+    keys = {'another-tenant-key'}
+    store = FakeStore()
+    store.batch_is_exist = lambda batch: [int(key in keys) for key in batch]
+    hashes = build_prefix_block_hashes(range(128), 64)
+    old_keys = set()
+    for generation in ['', 'first-wakeup', 'second-wakeup']:
+        worker, _ = make_worker(tmp_path, store=store, cache_generation=generation)
+        keys.update(old_keys)
+        assert worker.lookup(128, hashes) == 0
+        current = {build_store_key(worker.key_metadata, 0, h) for h in hashes}
+        assert not current.intersection(old_keys)
+        keys.update(current)
+        assert worker.lookup(128, hashes) == 128
+        worker.shutdown()
+        old_keys.update(current)
+    assert old_keys <= keys
+    assert 'another-tenant-key' in keys

From 682725874d2a7788ad6bc8d99f4e90f9085b7b34 Mon Sep 17 00:00:00 2001
From: RunningLeon <mnsheng@yeah.net>
Date: Wed, 21 Oct 2026 02:36:30 +0000
Subject: [PATCH 3/3] Revert "fix(pytorch): invalidate remote Store cache on
 sleep"

This reverts commit 3b189d07ee505ba280cb24910a7f89ff423ff565.
---
 lmdeploy/pytorch/engine/model_agent/agent.py  | 19 -------
 lmdeploy/pytorch/kv_connector/factory.py      |  2 -
 .../kv_connector/mooncake/store/connector.py  |  2 -
 .../kv_connector/mooncake/store/data.py       |  6 +--
 .../kv_connector/mooncake/store/worker.py     |  2 -
 lmdeploy/pytorch/paging/scheduler.py          |  4 +-
 .../engine/test_model_agent_kv_connector.py   | 49 ++-----------------
 .../test_mooncake_store_connector.py          |  2 -
 .../test_mooncake_store_worker.py             | 43 +++-------------
 9 files changed, 12 insertions(+), 117 deletions(-)

diff --git a/lmdeploy/pytorch/engine/model_agent/agent.py b/lmdeploy/pytorch/engine/model_agent/agent.py
index 6dfe12f344..28f31e01d6 100644
--- a/lmdeploy/pytorch/engine/model_agent/agent.py
+++ b/lmdeploy/pytorch/engine/model_agent/agent.py
@@ -10,7 +10,6 @@
 from multiprocessing.reduction import ForkingPickler
 from os import getenv
 from typing import Any
-from uuid import uuid4
 
 import numpy as np
 import pybase64
@@ -385,7 +384,6 @@ def __init__(
         self._cache_plan_block_nbytes: tuple[int, int]
         self.state_cache_engine = None
         self.kv_connector = None
-        self._kv_cache_generation = ''
         self.profiler: AgentProfiler = None
         try:
             self.guided_decoding_manager = GuidedDecodingManager(self.tokenizer, model_config.vocab_size)
@@ -1470,19 +1468,6 @@ def shutdown_kv_connector(self):
         self._release_completed_h2d_transfers()
         self._shutdown_kv_connector()
 
-    def _advance_kv_cache_generation(self):
-        """Isolate post-sleep remote KV from any previous weights or writes.
-
-        All DP/TP workers use the same fresh namespace. Keep it on the agent so rebuilding models or connectors during
-        staged wakeup cannot lose it. Old Store objects are left for eviction, not deleted from shared storage.
-        """
-        generation = [uuid4().hex if self.rank == 0 else None]
-        if self.dist_config.world_size > 1:
-            dist.broadcast_object_list(generation, src=0, group=self.dist_ctx.cpu_group)
-        self._kv_cache_generation = generation[0]
-        logger.info('Reset external KV cache generation: rank=%d generation=%s',
-                    self.rank, self._kv_cache_generation)
-
     def _shutdown_kv_connector(self):
         """Shutdown the connector before releasing its registered caches."""
         if self.kv_connector is None:
@@ -1560,7 +1545,6 @@ def build_cache_engine(self):
                 tp_rank=tp_rank,
                 tp_size=tp,
                 kv_head_replica_num=self.model_config.num_replicate_key_value_heads,
-                cache_generation=self._kv_cache_generation,
             )
             if self.kv_connector is not None:
                 self.kv_connector.register_kv_caches(
@@ -1923,10 +1907,7 @@ async def sleep(self, level: int = 1):
         device = 'cpu' if level == 1 else 'meta'
         # Stop producers of GPU work and wait for queued work before closing the
         # connector threads that still reference registered cache addresses.
-        had_kv_connector = self.kv_connector is not None
         self.shutdown_kv_connector()
-        if had_kv_connector:
-            self._advance_kv_cache_generation()
         self.cache_engine = None
         self.state_cache_engine = None
         self.reset_graph_runner()
diff --git a/lmdeploy/pytorch/kv_connector/factory.py b/lmdeploy/pytorch/kv_connector/factory.py
index c93b33fcdc..8c9da3b1d5 100644
--- a/lmdeploy/pytorch/kv_connector/factory.py
+++ b/lmdeploy/pytorch/kv_connector/factory.py
@@ -74,7 +74,6 @@ def build_kv_connector(
     tp_rank: int = 0,
     tp_size: int = 1,
     kv_head_replica_num: int = 1,
-    cache_generation: str = '',
 ) -> KVConnectorBase | None:
     """Build the configured connector without importing optional backends
     eagerly."""
@@ -94,7 +93,6 @@ def build_kv_connector(
             tp_rank=tp_rank,
             tp_size=tp_size,
             kv_head_replica_num=kv_head_replica_num,
-            cache_generation=cache_generation,
         )
 
     raise ValueError(f'Unsupported KV connector: {connector_name!r}')
diff --git a/lmdeploy/pytorch/kv_connector/mooncake/store/connector.py b/lmdeploy/pytorch/kv_connector/mooncake/store/connector.py
index 49947c375a..7ad5466f24 100644
--- a/lmdeploy/pytorch/kv_connector/mooncake/store/connector.py
+++ b/lmdeploy/pytorch/kv_connector/mooncake/store/connector.py
@@ -38,7 +38,6 @@ def __init__(
         tp_rank: int = 0,
         tp_size: int = 1,
         kv_head_replica_num: int = 1,
-        cache_generation: str = '',
     ) -> None:
         super().__init__(role)
 
@@ -63,7 +62,6 @@ def __init__(
                 tp_rank=tp_rank,
                 tp_size=tp_size,
                 kv_head_replica_num=kv_head_replica_num,
-                cache_generation=cache_generation,
             )
 
     def _require_scheduler(self) -> MooncakeStoreScheduler:
diff --git a/lmdeploy/pytorch/kv_connector/mooncake/store/data.py b/lmdeploy/pytorch/kv_connector/mooncake/store/data.py
index 623c3dc8a0..e2d144b9e5 100644
--- a/lmdeploy/pytorch/kv_connector/mooncake/store/data.py
+++ b/lmdeploy/pytorch/kv_connector/mooncake/store/data.py
@@ -177,7 +177,6 @@ class MooncakeStoreKeyMetadata:
     tp_size: int
     block_size: int
     kv_head_replica_num: int = 1
-    cache_generation: str = ''
 
     def __post_init__(self) -> None:
         validate_kv_head_replica_num(self.kv_head_replica_num, self.tp_size)
@@ -201,11 +200,8 @@ def build_store_key(
         raise ValueError(f'block_hash must contain {MOONCAKE_BLOCK_HASH_BYTES} bytes')
 
     prefix = f'{metadata.cache_prefix}@' if metadata.cache_prefix else ''
-    # Preserve the initial shared namespace, but never reuse pre-sleep KV after
-    # wakeup. Lookup, load and save all construct keys through this function.
-    generation = f'@generation:{metadata.cache_generation}' if metadata.cache_generation else ''
     return (
-        f'{prefix}{metadata.model_name}{generation}'
+        f'{prefix}{metadata.model_name}'
         f'@tp_rank:{kv_head_rank}'
         f'@group:{group_id}'
         f'@{block_hash.hex()}'
diff --git a/lmdeploy/pytorch/kv_connector/mooncake/store/worker.py b/lmdeploy/pytorch/kv_connector/mooncake/store/worker.py
index ff6e3532ad..26846e2759 100644
--- a/lmdeploy/pytorch/kv_connector/mooncake/store/worker.py
+++ b/lmdeploy/pytorch/kv_connector/mooncake/store/worker.py
@@ -72,7 +72,6 @@ def __init__(
         tp_rank: int = 0,
         tp_size: int = 1,
         kv_head_replica_num: int = 1,
-        cache_generation: str = '',
         store_factory: StoreFactory | None = None,
         replicate_config: Any = None,
     ) -> None:
@@ -111,7 +110,6 @@ def __init__(
             tp_size=tp_size,
             block_size=cache_config.block_size,
             kv_head_replica_num=kv_head_replica_num,
-            cache_generation=cache_generation,
         )
 
         config_path = extra_config.get('mooncake_config_path')
diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py
index 6987cab607..2d1de83f8e 100644
--- a/lmdeploy/pytorch/paging/scheduler.py
+++ b/lmdeploy/pytorch/paging/scheduler.py
@@ -510,8 +510,8 @@ def finish_kv_transfers_after_worker_drain(self) -> None:
     def reset_cache_after_worker_drain(self) -> None:
         """Invalidate local cache identity after sleep discards worker storage.
 
-        Ended external loads and save leases must release their references first. Workers have already advanced the
-        external cache generation, so neither local nor remote pre-sleep KV can be reused after wakeup.
+        Ended external loads and save leases must release their references first. Remote Store objects are not
+        invalidated; future requests may restore them through the normal lookup/load completion protocol.
         """
         self.finish_kv_transfers_after_worker_drain()
         if self.sessions:
diff --git a/tests/pytorch/engine/test_model_agent_kv_connector.py b/tests/pytorch/engine/test_model_agent_kv_connector.py
index 25b73d80b8..89c3a8b8c7 100644
--- a/tests/pytorch/engine/test_model_agent_kv_connector.py
+++ b/tests/pytorch/engine/test_model_agent_kv_connector.py
@@ -79,7 +79,6 @@ def fake_connector(role, cache_config, **kwargs):
         'tp_rank': 3,
         'tp_size': 8,
         'kv_head_replica_num': 1,
-        'cache_generation': '',
     }
 
 
@@ -89,7 +88,6 @@ def _bare_model_agent():
     agent = BaseModelAgent.__new__(BaseModelAgent)
     agent.all_context = nullcontext
     agent.cache_config = _enabled_config()
-    agent._kv_cache_generation = 'post-sleep-generation'
     agent.model_config = SimpleNamespace(num_replicate_key_value_heads=4)
     agent.rank = 7
     agent.cache_stream = object()
@@ -174,7 +172,6 @@ def fake_build_connector(role, cache_config, **kwargs):
         'tp_rank': 3,
         'tp_size': 8,
         'kv_head_replica_num': 4,
-        'cache_generation': 'post-sleep-generation',
     }
     assert events[5] == (
         'register', {
@@ -328,8 +325,7 @@ def test_build_cache_engine_rejects_mismatched_mtp_parallel_config(monkeypatch):
     assert agent.kv_connector is None
 
 
-@pytest.mark.parametrize('level', [1, 2])
-def test_sleep_shuts_down_connector_before_dropping_cache(monkeypatch, level):
+def test_sleep_shuts_down_connector_before_dropping_cache(monkeypatch):
     from lmdeploy.pytorch.engine.model_agent import agent as agent_module
 
     events = []
@@ -345,8 +341,7 @@ def shutdown(self):
     agent.kv_connector = _Connector()
     agent.cache_engine = cache_engine
     agent.state_cache_engine = object()
-    agent.rank = 0
-    agent.dist_config = SimpleNamespace(dp=1, world_size=1)
+    agent.dist_config = SimpleNamespace(dp=1)
     agent.state = SimpleNamespace(is_sleeping=False,
                                   to_sleep=SimpleNamespace(clear=lambda: events.append('sleep-clear')))
     model = SimpleNamespace(to=lambda **kwargs: events.append('model-to'))
@@ -362,22 +357,13 @@ def shutdown(self):
     monkeypatch.setattr(agent_module.torch.cuda, 'synchronize', lambda: events.append('cuda-sync'))
     monkeypatch.setattr(agent_module.torch.cuda, 'empty_cache', lambda: events.append('empty-cache'))
 
-    def new_generation():
-        assert events[-1] == 'shutdown'
-        assert agent.cache_engine is cache_engine
-        events.append('new-generation')
-        return SimpleNamespace(hex='fresh-generation')
-
-    monkeypatch.setattr(agent_module, 'uuid4', new_generation)
-    asyncio.run(agent.sleep(level))
-    assert agent._kv_cache_generation == 'fresh-generation'
+    asyncio.run(agent.sleep())
 
     assert events == [
         'drain',
         'cuda-sync',
         'release-h2d',
         'shutdown',
-        'new-generation',
         'reset-graph',
         'model-to',
         'spec-model-to',
@@ -737,32 +723,3 @@ def test_wakeup_rebuilds_connector_with_kv_cache():
 
     assert events == ['build-cache', 'warmup']
     assert agent.state.is_sleeping is False
-
-
-def test_cache_generation_is_shared_across_workers_and_changes_each_cycle(monkeypatch):
-    from lmdeploy.pytorch.engine.model_agent import agent as agent_module
-
-    group = object()
-    shared = []
-    agents = [_bare_model_agent() for _ in range(4)]
-    for rank, agent in enumerate(agents):
-        agent.rank = rank
-        agent.dist_config = SimpleNamespace(world_size=4)
-        agent.dist_ctx = SimpleNamespace(cpu_group=group)
-
-    def broadcast(values, src, group):
-        assert src == 0 and group is agents[0].dist_ctx.cpu_group
-        if values[0] is not None:
-            shared[:] = values
-        else:
-            values[:] = shared
-
-    monkeypatch.setattr(agent_module.dist, 'broadcast_object_list', broadcast)
-    generations = set()
-    for _ in range(3):
-        for agent in agents:
-            agent._advance_kv_cache_generation()
-        generation = agents[0]._kv_cache_generation
-        assert generation and generation not in generations
-        assert all(agent._kv_cache_generation == generation for agent in agents)
-        generations.add(generation)
diff --git a/tests/pytorch/kv_connector/test_mooncake_store_connector.py b/tests/pytorch/kv_connector/test_mooncake_store_connector.py
index 8298104674..a13846f1d5 100644
--- a/tests/pytorch/kv_connector/test_mooncake_store_connector.py
+++ b/tests/pytorch/kv_connector/test_mooncake_store_connector.py
@@ -74,7 +74,6 @@ def test_connector_constructs_only_worker_delegate(cache_config):
         global_rank=7,
         tp_rank=3,
         tp_size=8,
-        cache_generation='after-sleep',
     )
 
     assert connector.connector_scheduler is None
@@ -84,7 +83,6 @@ def test_connector_constructs_only_worker_delegate(cache_config):
     assert connector.connector_worker.tp_rank == 3
     assert connector.connector_worker.tp_size == 8
     assert connector.kv_role == 'kv_both'
-    assert connector.connector_worker.key_metadata.cache_generation == 'after-sleep'
 
 
 def test_connector_requires_enabled_kv_transfer_config():
diff --git a/tests/pytorch/kv_connector/test_mooncake_store_worker.py b/tests/pytorch/kv_connector/test_mooncake_store_worker.py
index 77b2e7adac..7bd4597335 100644
--- a/tests/pytorch/kv_connector/test_mooncake_store_worker.py
+++ b/tests/pytorch/kv_connector/test_mooncake_store_worker.py
@@ -785,8 +785,7 @@ def test_registration_rejects_non_cuda_noncontiguous_and_empty_rows(tmp_path):
     worker.shutdown()
 
 
-@pytest.mark.parametrize('cache_generation', ['', 'after-sleep'])
-def test_async_load_writes_allocated_blocks_and_reports_partial_failure(tmp_path, cache_generation):
+def test_async_load_writes_allocated_blocks_and_reports_partial_failure(tmp_path):
     path = write_store_config(tmp_path)
     store = FakeStore(get_results=[0, -1])
     worker = MooncakeStoreWorker(
@@ -795,7 +794,6 @@ def test_async_load_writes_allocated_blocks_and_reports_partial_failure(tmp_path
         tp_rank=5,
         tp_size=8,
         kv_head_replica_num=4,
-        cache_generation=cache_generation,
         store_factory=lambda: store,
     )
     worker.register_kv_caches({
@@ -878,12 +876,8 @@ def test_async_load_writes_target_and_mtp_cache_rows(tmp_path, monkeypatch):
 
 
 @pytest.mark.parametrize('fa_exists', [False, True])
-@pytest.mark.parametrize('cache_generation', ['', 'after-sleep'])
-def test_hybrid_save_keeps_group_zero_format_and_writes_every_state_rank(
-        tmp_path, monkeypatch, fa_exists, cache_generation):
-    metadata = MooncakeStoreKeyMetadata('model', 'prefix', 8, 64, kv_head_replica_num=4,
-                                        cache_generation=cache_generation)
-    namespace = 'prefix@model' + (f'@generation:{cache_generation}' if cache_generation else '')
+def test_hybrid_save_keeps_group_zero_format_and_writes_every_state_rank(tmp_path, monkeypatch, fa_exists):
+    metadata = MooncakeStoreKeyMetadata('model', 'prefix', 8, 64, kv_head_replica_num=4)
     hashes = build_prefix_block_hashes(range(512), 64)
     request = MooncakeStoreSaveRequest(1, 2, 0, tuple(range(8)), tuple(range(8)), hashes,
                                       state=MooncakeStoreStateSave(1, 3, 512))
@@ -896,7 +890,7 @@ def test_hybrid_save_keeps_group_zero_format_and_writes_every_state_rank(
             row_block_sizes=(64,), num_gpu_blocks=8, key_metadata=metadata, global_rank=rank, tp_rank=rank,
             tp_size=8, completion_callback=completed.append, replicate_config=object(),
             state_registrations=(MooncakeStoreStateRegistration('state', 0x2000, 64, 4, 16),))
-        state_key = f'{namespace}@tp_rank:{rank}@group:1@{hashes[-1].hex()}'
+        state_key = f'prefix@model@tp_rank:{rank}@group:1@{hashes[-1].hex()}'
         entries = sender._owned_entries(request)
         # Exercise both orders: a state entry need not follow the FA entries.
         if rank % 2:
@@ -914,7 +908,7 @@ def test_hybrid_save_keeps_group_zero_format_and_writes_every_state_rank(
         written.extend(keys)
     expected_states = {build_store_key(metadata, rank, hashes[-1], group_id=1) for rank in range(8)}
     expected_fa = set() if fa_exists else {
-        f'{namespace}@tp_rank:{rank}@group:0@{block_hash.hex()}' for rank in range(2) for block_hash in hashes}
+        f'prefix@model@tp_rank:{rank}@group:0@{block_hash.hex()}' for rank in range(2) for block_hash in hashes}
     assert len(written) == len(expected_states | expected_fa)
     assert set(written) == expected_states | expected_fa
 
@@ -924,7 +918,7 @@ def test_hybrid_save_keeps_group_zero_format_and_writes_every_state_rank(
     config = replace(make_cache_config(write_store_config(tmp_path)),
                      states_shapes=[((2,), torch.float32)], prefill_save_alignment=512)
     config.kv_transfer_config.kv_connector_extra_config.update(model_name='model', cache_prefix='prefix')
-    worker = MooncakeStoreWorker(config, tp_size=8, kv_head_replica_num=4, cache_generation=cache_generation,
+    worker = MooncakeStoreWorker(config, tp_size=8, kv_head_replica_num=4,
                                 store_factory=lambda: FakeStore(existing_keys=existing_keys))
     assert worker.lookup(512, hashes) == 512
     worker.shutdown()
@@ -974,11 +968,9 @@ def read_snapshot():
     assert store.put_calls[0][2] == [[6, 5, 5]]
 
 
-@pytest.mark.parametrize('cache_generation', ['', 'after-sleep'])
 def test_async_save_waits_for_forward_and_writes_only_owned_missing_blocks(
     tmp_path,
     monkeypatch,
-    cache_generation,
 ):
     path = write_store_config(tmp_path)
     event = FakeCudaEvent()
@@ -998,7 +990,6 @@ def assert_forward_is_ready():
         tp_rank=3,
         tp_size=4,
         kv_head_replica_num=2,
-        cache_generation=cache_generation,
         store_factory=lambda: store,
         replicate_config=replicate_config,
     )
@@ -1095,28 +1086,6 @@ def test_shutdown_is_idempotent_and_propagates_close_error(tmp_path):
     assert worker.store is None
 
 
-def test_cache_generation_invalidates_remote_hits_without_deleting_shared_objects(tmp_path):
-    # Store storage outlives workers. Even delayed writes to an old generation
-    # cannot be mistaken for KV produced after a sleep/weight update.
-    keys = {'another-tenant-key'}
-    store = FakeStore()
-    store.batch_is_exist = lambda batch: [int(key in keys) for key in batch]
-    hashes = build_prefix_block_hashes(range(128), 64)
-    old_keys = set()
-    for generation in ['', 'first-wakeup', 'second-wakeup']:
-        worker, _ = make_worker(tmp_path, store=store, cache_generation=generation)
-        keys.update(old_keys)
-        assert worker.lookup(128, hashes) == 0
-        current = {build_store_key(worker.key_metadata, 0, h) for h in hashes}
-        assert not current.intersection(old_keys)
-        keys.update(current)
-        assert worker.lookup(128, hashes) == 128
-        worker.shutdown()
-        old_keys.update(current)
-    assert old_keys <= keys
-    assert 'another-tenant-key' in keys
-
-
 class BytesStore(FakeStore):
 
     def __init__(self):