diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 9aae15d6d..d82141228 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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: diff --git a/python/sglang/srt/mem_cache/evict_policy.py b/python/sglang/srt/mem_cache/evict_policy.py index 30ab1983f..2db247df4 100644 --- a/python/sglang/srt/mem_cache/evict_policy.py +++ b/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) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 3e6059c7c..87a2488c5 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/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( diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index b4598aad9..df824ac2c 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/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) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 56e0b2c73..68709d69f 100644 --- a/python/sglang/srt/server_args.py +++ b/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" diff --git a/test/registered/unit/mem_cache/test_evict_policy.py b/test/registered/unit/mem_cache/test_evict_policy.py index 0dd2eed5d..11650b5cc 100644 --- a/test/registered/unit/mem_cache/test_evict_policy.py +++ b/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() diff --git a/test/registered/unit/mem_cache/test_radix_cache_unit.py b/test/registered/unit/mem_cache/test_radix_cache_unit.py index ae01c8eb6..fa360906b 100644 --- a/test/registered/unit/mem_cache/test_radix_cache_unit.py +++ b/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):