CP HiCache: replicated-clock SLRU eviction, collective-safe (Phase 1)

Make CP shared-KV HiCache eviction collective-free and scan-resistant by
removing every per-rank wall-clock input from the eviction/admission decisions
(any such input desyncs the must-be-replicated victim set / batch across the 8
CP ranks and can deadlock the collective-coupled writeback/load-back).

- Replicated logical clock: TreeNode.next_access_time() (a process-global
  monotonic counter, mirrors mamba/swa's get_last_access_time) replaces
  time.monotonic() as the source of last_access_time at every radix bump site
  (radix_cache __init__/match/insert; hiradix CP match/insert/_insert_helper_host;
  reset() zeroes it). creation_time/pin_expiry intentionally stay wall-clock.
  Because match/insert events are replicated (reqs broadcast from rank 0 over the
  replicated tree), last_access_time is now identical on every rank. Unique ints
  also give a strict total order (no LRU tie ambiguity). No duration arithmetic
  reads last_access_time, so non-CP LRU is unchanged; mamba/swa use their own
  TreeNode and are untouched.
- CpReplicatedSLRUStrategy = (is_protected[hit>=2], last_access_time, node.id):
  scan-resistant (cold one-shots stay probationary, evicted before reused/
  protected prefixes), recency-within-segment ages out cold (not LFU), node.id
  is the deterministic total-order tiebreak (heaps never compare TreeNodes).
  Overridden for CP so it reaches every get_priority eviction site; the host
  write-admission key _cp_host_evict_key is switched from FIFO-by-id to it.
- Co-fixes for the same per-rank-wall-clock class: pin_prefix is forbidden under
  CP (its pin_expiry victim-eligibility check is per-rank wall-clock; SLRU
  scan-resistance covers the benefit); the affinity head_age_s admission input is
  disabled (=0.0, relying on the replicated head_defer_count bound); and a
  server_args fail-fast guards CP shared-KV against SGLANG_REQ_WAITING_TIMEOUT>0
  (its waiting-queue pruning is per-rank wall-clock).

Tests: test_evict_policy.py adds TestCpReplicatedSLRUStrategy (scan-resistance,
recency, total-order, full ordering) -- 30 passed in the dev-cu13 container;
test_radix_cache_unit.py adds logical-clock determinism/total-order tests.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

python/sglang/srt/managers/scheduler.py

@@ -2665,12 +2665,16 @@ class Scheduler(
                     batch_warm_led=bool(affinity_batch_warm_led),
                     is_head=is_head,
                     head_defer_count=req.affinity_defer_count,
-                    head_age_s=(
-                        time.perf_counter()
-                        - req.time_stats.wait_queue_entry_time
-                        if req.time_stats.wait_queue_entry_time > 0.0
-                        else 0.0
-                    ),
+                    # head_age_s DISABLED (=0.0) for CP collective-safety:
+                    # time.perf_counter() is per-rank, so at the MAX_AGE_S boundary
+                    # ranks could admit DIFFERENT batches → permanent node.id /
+                    # access-clock / eviction desync (breaks the replicated
+                    # event-stream premise). The replicated head_defer_count
+                    # (MAX_DEFER passes) already bounds cold-head starvation, and
+                    # the SKIP_COLD path only triggers under load (frequent passes).
+                    # Restore a time-like bound only via a replicated logical age
+                    # (ticks-since-enqueue). MAX_AGE_S is currently unused.
+                    head_age_s=0.0,
                     window_used=affinity_window_used,
                 )
                 if decision is AffinityDecision.SKIP_COLD:

python/sglang/srt/mem_cache/evict_policy.py

@@ -63,3 +63,35 @@ class SLRUStrategy(EvictionStrategy):
 
         is_protected = 1 if node.hit_count >= self.protected_threshold else 0
         return (is_protected, node.last_access_time)
+
+
+class CpReplicatedSLRUStrategy(EvictionStrategy):
+    """SLRU for CP shared-KV HiCache, made collective-safe by a strict total order.
+
+    Value = (is_protected, last_access_time, node.id), smallest evicted first:
+      - is_protected (hit_count >= threshold): probationary nodes (cold one-shots,
+        e.g. a long-request injector) are evicted before ANY protected node →
+        scan-resistance for the shared working set.
+      - last_access_time: LRU within a segment, so a now-cold protected prefix
+        ages out by recency (NOT by frequency — avoids LFU pinning a
+        stale-high-hit-count prefix).
+      - node.id: final deterministic tiebreak → a strict TOTAL order (the
+        per-call logical clock can repeat a value across nodes touched in one
+        match; node.id disambiguates).
+
+    Every field is rank-replicated: last_access_time is the logical clock
+    (TreeNode.next_access_time()), node.id/hit_count advance identically on all
+    CP ranks. So all ranks select identical victims with no collective. (Plain
+    LRU/LFU/SLRU here tiebreak on per-rank wall-clock last_access_time and would
+    diverge victims → collective deadlock.)
+    """
+
+    def __init__(self, protected_threshold: int = 2):
+        # protected_threshold T=2: a one-shot (hit_count==1) stays probationary;
+        # a genuinely-reused prefix (hit>=2) is protected. T is the scan-resistance
+        # tuning knob (design item #3, to be measured).
+        self.protected_threshold = protected_threshold
+
+    def get_priority(self, node: "TreeNode") -> Tuple[int, float, int]:
+        is_protected = 1 if node.hit_count >= self.protected_threshold else 0
+        return (is_protected, node.last_access_time, node.id)

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -671,6 +671,20 @@ class HiRadixCache(RadixCache):
 
         super().__init__(params=params)
 
+        if self._uses_cp_hicache:
+            # CP shared-KV eviction must be COLLECTIVE-FREE: every rank must pick
+            # identical victims from rank-replicated state, else partial backups /
+            # all-reduce-shape deadlock. The base LRU/SLRU strategies tiebreak on
+            # per-rank wall-clock last_access_time and diverge. Override with the
+            # replicated-clock SLRU value (is_protected, last_access_time[logical],
+            # node.id) — a strict total order over replicated state. Reaches every
+            # eviction key site that uses eviction_strategy.get_priority (device
+            # owner-lane score, host physical-slot heap, generic evict); the host
+            # write-admission key _cp_host_evict_key is switched to it explicitly.
+            from sglang.srt.mem_cache.evict_policy import CpReplicatedSLRUStrategy
+
+            self.eviction_strategy = CpReplicatedSLRUStrategy()
+
     def _cp_hicache_all_reduce(self, tensor, *, op, group, tag: str) -> None:
         start = time.perf_counter()
         torch.distributed.all_reduce(tensor, op=op, group=group)
@@ -922,10 +936,11 @@ class HiRadixCache(RadixCache):
         )
 
     def _cp_host_evict_key(self, node: TreeNode):
-        return (
-            int(getattr(node, "priority", 0)),
-            int(getattr(node, "id", 0)),
-        )
+        # Replicated-clock SLRU total order (is_protected, last_access_time[logical],
+        # node.id) — replaces the old FIFO-by-node.id key. Same strategy the device
+        # owner-lane / physical-slot eviction paths use, so all host/device victim
+        # selection is consistent and collective-free.
+        return self.eviction_strategy.get_priority(node)
 
     def _cp_host_leaf_is_plannable_victim(self, node: TreeNode) -> bool:
         if node == getattr(self, "root_node", None):
@@ -1249,6 +1264,10 @@ class HiRadixCache(RadixCache):
                     )
                 if contribution <= 0 and unlock_contribution <= 0:
                     continue
+                # get_priority returns the CpReplicatedSLRUStrategy tuple
+                # (is_protected, last_access_time[logical], node.id) under CP — a
+                # rank-replicated value; the trailing node.id keeps the full score a
+                # strict total order even if the strategy is ever swapped.
                 score = (
                     -int(contribution),
                     -int(unlock_contribution),
@@ -1996,6 +2015,7 @@ class HiRadixCache(RadixCache):
 
     def reset(self):
         TreeNode.counter = 0
+        TreeNode.last_access_counter = 0
         self.cache_controller.reset()
         self.token_to_kv_pool_host.clear()
         self.cache_controller.clear_draft_host_pool()
@@ -3149,6 +3169,14 @@ class HiRadixCache(RadixCache):
         self, token_ids: List[int], ttl_seconds: int = 300
     ) -> Tuple[int, Optional[str]]:
         """Pin nodes along a prefix path. Returns (nodes_pinned, reject_reason)."""
+        if self._uses_cp_hicache:
+            # Forbidden under CP shared-KV: pin_expiry is a per-rank wall-clock TTL
+            # and the eviction victim-eligibility check (time.monotonic() <=
+            # pin_expiry) would then diverge across CP ranks → non-replicated
+            # victim set → collective-free eviction breaks. SLRU scan-resistance
+            # already protects the hot working set; re-enable only with a
+            # replicated logical TTL.
+            return (0, "pin_prefix is not supported under CP shared-KV HiCache")
         if self.disable or not token_ids:
             return (0, None)
 
@@ -4176,7 +4204,11 @@ class HiRadixCache(RadixCache):
     def _insert_helper_host(
         self, node: TreeNode, key: RadixKey, host_value, hash_value
     ):
-        node.last_access_time = time.monotonic()
+        # NOTE: dead under CP today (storage forbidden); converted for whole-repo
+        # correctness so a future CP-aware L3 prefetch insert stays rank-replicated.
+        # (The prefetch-completion insert is MIN-synchronized, so the bump order is
+        # replicated; if it ever ran at a non-replicated point this must be excluded.)
+        node.last_access_time = TreeNode.next_access_time()
         if len(key) == 0:
             return 0
 
@@ -4185,7 +4217,7 @@ class HiRadixCache(RadixCache):
         matched_length = 0
         while len(key) > 0 and child_key in node.children.keys():
             node = node.children[child_key]
-            node.last_access_time = time.monotonic()
+            node.last_access_time = TreeNode.next_access_time()
             # Refresh pin TTL on host insert hit
             if self._is_pinned(node):
                 node.pin_expiry = time.monotonic() + node.pin_ttl
@@ -4219,13 +4251,13 @@ class HiRadixCache(RadixCache):
     def _match_prefix_helper(
         self, node: TreeNode, key: RadixKey, *, floor_exact_key: bool = True
     ):
-        node.last_access_time = time.monotonic()
+        node.last_access_time = TreeNode.next_access_time()
         child_key = self.get_child_key_fn(key)
         value = []
 
         while len(key) > 0 and child_key in node.children.keys():
             child = node.children[child_key]
-            child.last_access_time = time.monotonic()
+            child.last_access_time = TreeNode.next_access_time()
             # Refresh pin TTL on cache hit
             if self._is_pinned(child):
                 child.pin_expiry = time.monotonic() + child.pin_ttl
@@ -4351,7 +4383,7 @@ class HiRadixCache(RadixCache):
         while len(key) > 0 and child_key in node.children.keys():
             parent_node = node
             node = node.children[child_key]
-            node.last_access_time = time.monotonic()
+            node.last_access_time = TreeNode.next_access_time()
             node.priority = max(node.priority, priority)
             raw_prefix_len = self.key_match_fn(node.key, key)
             prefix_len = self._cp_floor_exact_valid_tail_extension_len(

python/sglang/srt/mem_cache/radix_cache.py

@@ -186,6 +186,23 @@ def floor_to_page_len(length: int, page_size: int) -> int:
 class TreeNode:
 
     counter = 0
+    # Replicated logical access clock (mirrors mamba_radix_cache's
+    # get_last_access_time pattern). A process-global monotonically-increasing
+    # counter used as the source of ``last_access_time``. Under CP shared-KV
+    # every rank advances it through the IDENTICAL sequence of match/insert
+    # events (requests are broadcast from rank 0 and the radix tree is
+    # replicated), so ``last_access_time`` is rank-replicated and eviction
+    # victim order is collective-free. The old per-rank ``time.monotonic()``
+    # diverged victim order across ranks and could deadlock the collective-
+    # coupled writeback/load-back. Unique integers also give a strict total
+    # order (no tie ambiguity). Reset alongside ``counter`` in HiRadixCache.reset.
+    last_access_counter = 0
+
+    @staticmethod
+    def next_access_time() -> int:
+        v = TreeNode.last_access_counter
+        TreeNode.last_access_counter += 1
+        return v
 
     def __init__(self, id: Optional[int] = None, priority: int = 0):
         self.children = {}
@@ -197,7 +214,11 @@ class TreeNode:
             0.0  # absolute expiry time (time.monotonic()), 0 = not pinned
         )
         self.pin_ttl: int = 0  # original TTL in seconds, for refresh-on-hit
-        self.last_access_time = time.monotonic()
+        self.last_access_time = TreeNode.next_access_time()
+        # creation_time stays wall-clock: it is only read by FIFO/FILO eviction
+        # (not the CP/default LRU/SLRU path). FIFO/FILO are not CP-safe today and
+        # are not the default; if ever enabled under CP, route this to the
+        # logical clock too.
         self.creation_time = time.monotonic()
 
         self.hit_count = 0
@@ -926,7 +947,7 @@ class RadixCache(BasePrefixCache):
     ##### Internal Helper Functions #####
 
     def _match_prefix_helper(self, node: TreeNode, key: RadixKey):
-        access_time = time.monotonic()
+        access_time = TreeNode.next_access_time()
         node.last_access_time = access_time
 
         child_key = self.get_child_key_fn(key)
@@ -994,7 +1015,7 @@ class RadixCache(BasePrefixCache):
         # Convert None priority to 0
         if priority is None:
             priority = 0
-        access_time = time.monotonic()
+        access_time = TreeNode.next_access_time()
         node.last_access_time = access_time
         # Update priority along the path (take max to propagate higher priority)
         node.priority = max(node.priority, priority)

python/sglang/srt/server_args.py

@@ -994,6 +994,19 @@ class ServerArgs:
             "hicache_storage_backend in the host-only CP HiCache stage. "
             "Disable hicache_storage_backend or disable CP shared KV."
         )
+        assert not (
+            self.enable_nsa_prefill_cp_shared_kv
+            and envs.SGLANG_REQ_WAITING_TIMEOUT.get() > 0
+        ), (
+            "enable_nsa_prefill_cp_shared_kv is incompatible with "
+            "SGLANG_REQ_WAITING_TIMEOUT > 0: _abort_on_waiting_timeout prunes the "
+            "waiting queue using per-rank time.perf_counter(), so CP ranks could "
+            "drop different requests at the deadline boundary and desync the "
+            "(must-be-replicated) waiting queue -> radix node ids / eviction order "
+            "-> collective deadlock. Unset SGLANG_REQ_WAITING_TIMEOUT or disable CP "
+            "shared KV. (Same per-rank-wall-clock class as the disabled affinity "
+            "head_age_s and the prefetch-gating sync requirement.)"
+        )
         assert not (
             self.enable_nsa_prefill_cp_shared_kv
             and self.disaggregation_mode == "prefill"

test/registered/unit/mem_cache/test_evict_policy.py

@@ -8,6 +8,7 @@ import unittest
 from unittest.mock import MagicMock
 
 from sglang.srt.mem_cache.evict_policy import (
+    CpReplicatedSLRUStrategy,
     FIFOStrategy,
     FILOStrategy,
     LFUStrategy,
@@ -24,6 +25,7 @@ def _make_node(**kwargs):
     node.hit_count = kwargs.get("hit_count", 0)
     node.creation_time = kwargs.get("creation_time", 0.0)
     node.priority = kwargs.get("priority", 0)
+    node.id = kwargs.get("id", 0)
     return node
 
 
@@ -214,5 +216,74 @@ class TestEvictionOrdering(unittest.TestCase):
         self.assertEqual(actual, expected)
 
 
+class TestCpReplicatedSLRUStrategy(unittest.TestCase):
+    """CP shared-KV eviction strategy: scan-resistant SLRU + a strict total order.
+
+    Value = (is_protected, last_access_time, node.id). last_access_time is the
+    replicated logical clock under CP; node.id is the final tiebreak.
+    """
+
+    def setUp(self):
+        self.strategy = CpReplicatedSLRUStrategy(protected_threshold=2)
+
+    def test_priority_shape_probationary(self):
+        node = _make_node(hit_count=1, last_access_time=5.0, id=7)
+        self.assertEqual(self.strategy.get_priority(node), (0, 5.0, 7))
+
+    def test_priority_shape_protected(self):
+        node = _make_node(hit_count=2, last_access_time=5.0, id=3)
+        self.assertEqual(self.strategy.get_priority(node), (1, 5.0, 3))
+
+    def test_scan_resistance_probationary_before_protected(self):
+        # A cold one-shot (hit_count=1, even if just touched) is evicted before
+        # any reused/protected prefix (hit_count>=2, even if much older). This is
+        # what shields the hot working set from a cold long-request injector scan.
+        cold_oneshot = _make_node(hit_count=1, last_access_time=100.0, id=1)
+        hot_prefix = _make_node(hit_count=50, last_access_time=1.0, id=2)
+        self.assertLess(
+            self.strategy.get_priority(cold_oneshot),
+            self.strategy.get_priority(hot_prefix),
+        )
+
+    def test_recency_within_segment_ages_out_cold(self):
+        # Within the protected segment, the LESS-recent node is evicted first
+        # (recency, NOT frequency) -> a now-cold high-hit-count prefix ages out,
+        # which pure LFU would wrongly pin.
+        stale_high_freq = _make_node(hit_count=1000, last_access_time=1.0, id=1)
+        recent_low_freq = _make_node(hit_count=2, last_access_time=10.0, id=2)
+        self.assertLess(
+            self.strategy.get_priority(stale_high_freq),
+            self.strategy.get_priority(recent_low_freq),
+        )
+
+    def test_node_id_makes_strict_total_order(self):
+        # Nodes touched in one match can share a logical last_access_time; node.id
+        # breaks the tie so the value is a STRICT total order (heaps never fall
+        # through to comparing TreeNode objects -> deterministic across CP ranks).
+        a = _make_node(hit_count=1, last_access_time=5.0, id=3)
+        b = _make_node(hit_count=1, last_access_time=5.0, id=7)
+        self.assertNotEqual(
+            self.strategy.get_priority(a), self.strategy.get_priority(b)
+        )
+        self.assertLess(
+            self.strategy.get_priority(a), self.strategy.get_priority(b)
+        )
+
+    def test_default_threshold_is_2(self):
+        self.assertEqual(CpReplicatedSLRUStrategy().protected_threshold, 2)
+
+    def test_full_eviction_ordering(self):
+        nodes = [
+            _make_node(hit_count=5, last_access_time=1.0, id=10),  # protected, old
+            _make_node(hit_count=1, last_access_time=9.0, id=11),  # probationary, new
+            _make_node(hit_count=1, last_access_time=2.0, id=12),  # probationary, old
+            _make_node(hit_count=1, last_access_time=2.0, id=4),  # probationary, old, lower id
+            _make_node(hit_count=3, last_access_time=8.0, id=13),  # protected, new
+        ]
+        order = sorted(nodes, key=self.strategy.get_priority)
+        # probationary first (by last_access, then id), then protected (by last_access)
+        self.assertEqual([n.id for n in order], [4, 12, 11, 10, 13])
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_radix_cache_unit.py

@@ -17,12 +17,15 @@ Usage:
     python -m pytest test_radix_cache_unit.py::TestRadixCache::test_insert_basic
 """
 
-import sys
 import unittest.mock
 
-for _mod in ("sgl_kernel", "sgl_kernel.kvcacheio"):
-    if _mod not in sys.modules:
-        sys.modules[_mod] = unittest.mock.MagicMock()
+# NB: do NOT stub sgl_kernel here. This test is registered GPU-only
+# (register_cuda_ci/register_amd_ci), so the real sgl_kernel is always present,
+# and the import chain touches sgl_kernel solely through fp8_kernel.py under
+# `if _is_cuda:`. A MagicMock would satisfy `from sgl_kernel import ...` while
+# leaving the real C++ ops unregistered, so fp8_kernel's module-level
+# torch.library.register_fake("sgl_kernel::...") would raise
+# "operator sgl_kernel::... does not exist".
 
 from sglang.srt.mem_cache.common import available_and_evictable_str
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
@@ -235,12 +238,33 @@ class TestTreeNode(unittest.TestCase):
     def test_lt_comparison(self):
         """Test less than comparison based on last_access_time."""
         node1 = TreeNode()
-        time.sleep(0.001)  # Small delay to ensure different timestamps
+        time.sleep(0.001)  # (now a no-op: last_access_time is the logical clock)
         node2 = TreeNode()
 
         self.assertTrue(node1 < node2)
         self.assertFalse(node2 < node1)
 
+    def test_next_access_time_is_monotonic_logical_clock(self):
+        """next_access_time() is a replicated monotonic logical counter (NOT
+        wall-clock): consecutive calls strictly increase by 1 and return ints.
+        Under CP this is what makes last_access_time rank-replicated."""
+        a = TreeNode.next_access_time()
+        b = TreeNode.next_access_time()
+        c = TreeNode.next_access_time()
+        self.assertIsInstance(a, int)
+        self.assertEqual(b, a + 1)
+        self.assertEqual(c, b + 1)
+
+    def test_creation_order_is_strict_total_order(self):
+        """Nodes take last_access_time from the logical clock at construction, so
+        creation order is a strict total order with no wall-clock dependence or
+        ties (no sleep needed)."""
+        n1 = TreeNode()
+        n2 = TreeNode()
+        n3 = TreeNode()
+        self.assertLess(n1.last_access_time, n2.last_access_time)
+        self.assertLess(n2.last_access_time, n3.last_access_time)
+
 
 class TestTreeNodeChildren(CustomTestCase):
     def test_missing_child_access_does_not_create_orphan_node(self):