diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 3c27a7927..2b7926551 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -732,23 +732,25 @@ class HiRadixCache(RadixCache): device_indices=node.value, node_id=node.id, ) + metadata = getattr(result, "metadata", None) required_host_slots = 0 - if getattr(result, "metadata", None) is None: + if metadata is None: required_host_slots = result.required_host_slots self._evict_host_for_physical_slots( required_host_slots, synchronize_across_ranks=getattr(self, "tp_world_size", 1) > 1, ) - if getattr(result, "metadata", None) is None: + if metadata is None: result = self.cache_controller.write( device_indices=node.value, node_id=node.id, ) - if getattr(result, "metadata", None) is None: + metadata = getattr(result, "metadata", None) + if metadata is None: return 0 node.host_len = len(node.value) - node.cp_hicache = result.metadata + node.cp_hicache = metadata node.host_value = None self.ongoing_write_through[node.id] = node if not write_back: @@ -1100,7 +1102,9 @@ class HiRadixCache(RadixCache): self, required_host_slots: int, synchronize_across_ranks: bool = False ) -> int: synchronize_across_ranks = ( - synchronize_across_ranks and getattr(self, "tp_world_size", 1) > 1 + synchronize_across_ranks + and getattr(self, "tp_group", None) is not None + and getattr(self, "tp_world_size", 1) > 1 ) if required_host_slots <= 0 and not synchronize_across_ranks: return 0 @@ -1148,9 +1152,7 @@ class HiRadixCache(RadixCache): if x.host_ref_counter > 0: continue - host_indices = ( - x.cp_hicache.host_indices if x.cp_hicache is not None else None - ) + host_indices = self._node_host_evict_indices(x) physical_count = len(host_indices) if host_indices is not None else 0 self._record_remove_event(x) @@ -1197,7 +1199,9 @@ class HiRadixCache(RadixCache): # Block deleted entirely (GPU already evicted, now CPU freed) -- # emit BlockRemoved so the router removes this block from its index. self._record_remove_event(x) - num_evicted += self.cache_controller.evict_host(x.host_value) + num_evicted += self.cache_controller.evict_host( + self._node_host_evict_indices(x) + ) parent = self._remove_host_leaf(x) @@ -1581,8 +1585,10 @@ class HiRadixCache(RadixCache): while last_node.evicted: host_hit_length += len(last_node.host_value) last_node = last_node.parent - while not last_host_node.backuped: + while last_host_node != self.root_node and not last_host_node.backuped: last_host_node = last_host_node.parent + if not last_host_node.backuped: + last_host_node = self.root_node return MatchResult( device_indices=value, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 009570871..82e0f14e0 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -745,15 +745,15 @@ class ServerArgs: Orchestrates the handling of various server arguments, ensuring proper configuration and validation. """ - # Normalize load balancing defaults early (before dummy-model short-circuit). - self._handle_load_balance_method() - # Validate SSL arguments early (before dummy-model short-circuit). self._handle_ssl_validation() # Validate CP shared KV constraints early (before dummy-model short-circuit). self._handle_cp_shared_kv_validation() + # Normalize load balancing defaults early (before dummy-model short-circuit). + self._handle_load_balance_method() + if self.model_path.lower() in ["none", "dummy"]: # Skip for dummy models return diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index f91f2da43..6ef679017 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -416,6 +416,9 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase): cache._clear_pin = lambda node: None cache._record_remove_event = lambda node: None cache._update_host_leaf_status = lambda node: None + cache._node_host_evict_indices = lambda node: torch.tensor( + [99], dtype=torch.int64 + ) freed = [] cache.cache_controller = type( "Controller", @@ -441,7 +444,7 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase): physical_freed = cache._evict_host_for_physical_slots(1) self.assertEqual(physical_freed, 1) - self.assertEqual(freed[0].tolist(), [70]) + self.assertEqual(freed[0].tolist(), [99]) self.assertEqual(node.host_len, 0) self.assertIsNone(node.cp_hicache) @@ -567,7 +570,7 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True cache.tp_world_size = 2 - cache.tp_group = None + cache.tp_group = object() cache.root_node = TreeNode() cache.root_node.key = RadixKey([]) cache.evictable_host_leaves = set() @@ -600,6 +603,7 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase): cache.evictable_host_leaves.add(node) def mark_not_done(done_tensor, op=None, group=None): + self.assertIs(group, cache.tp_group) done_tensor.fill_(0) with patch("torch.distributed.all_reduce", side_effect=mark_not_done): @@ -613,6 +617,55 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase): self.assertEqual(node.host_len, 0) self.assertIsNone(node.cp_hicache) + def test_cp_host_eviction_skips_all_reduce_without_tp_group(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.tp_world_size = 2 + cache.tp_group = None + cache.root_node = TreeNode() + cache.root_node.key = RadixKey([]) + cache.evictable_host_leaves = set() + cache.get_child_key_fn = lambda key: key.token_ids[0] + cache.eviction_strategy = type( + "Strategy", (), {"get_priority": lambda self, node: 0} + )() + cache._clear_pin = lambda node: None + cache._record_remove_event = lambda node: None + cache.cache_controller = type( + "Controller", + (), + { + "evict_host": lambda self, indices: (_ for _ in ()).throw( + AssertionError("must not evict host slots") + ) + }, + )() + node = TreeNode() + node.parent = cache.root_node + node.key = RadixKey([1, 2, 3, 4]) + node.value = None + node.host_len = 4 + node.cp_hicache = CpHiCacheNodeMetadata( + logical_len=4, + owned_positions=torch.empty((0,), dtype=torch.int64), + host_indices=torch.empty((0,), dtype=torch.int64), + ) + cache.root_node.children[1] = node + cache.evictable_host_leaves.add(node) + + with patch( + "torch.distributed.all_reduce", + side_effect=AssertionError("must not all_reduce without tp_group"), + ): + physical_freed = cache._evict_host_for_physical_slots( + 0, synchronize_across_ranks=True + ) + + self.assertEqual(physical_freed, 0) + self.assertIn(1, cache.root_node.children) + self.assertEqual(node.host_len, 4) + self.assertIsNotNone(node.cp_hicache) + class TestHiRadixCacheCPLoadBack(CustomTestCase): def test_cp_load_back_uses_host_len_not_host_value(self): @@ -742,6 +795,35 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase): self.assertIs(result.last_host_node, cache.root_node) self.assertEqual(result.host_hit_length, 0) + def test_non_cp_match_prefix_uses_root_when_no_host_backup_exists(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = False + cache.device = "cpu" + cache.disable = False + cache.page_size = 1 + cache.get_child_key_fn = lambda key: key.token_ids[0] + cache.key_match_fn = lambda child_key, key: sum( + 1 for lhs, rhs in zip(child_key.token_ids, key.token_ids) if lhs == rhs + ) + cache.maybe_bigram_convert = lambda key: (key, None) + root = TreeNode() + root.key = RadixKey([]) + root.value = torch.empty((0,), dtype=torch.int64) + root.host_value = None + cache.root_node = root + node = TreeNode() + node.parent = root + node.key = RadixKey(list(range(4))) + node.value = torch.arange(4, dtype=torch.int64) + node.host_value = None + root.children[0] = node + + result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(4))))) + + self.assertEqual(result.host_hit_length, 0) + self.assertIs(result.last_host_node, root) + self.assertIs(result.last_device_node, node) + def test_cp_load_back_non_evicted_node_returns_none_without_loading(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index c9f884b20..0fafd503d 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -389,7 +389,7 @@ class TestHiCacheArgs(CustomTestCase): enable_nsa_prefill_context_parallel=True, enable_nsa_prefill_cp_shared_kv=True, nsa_prefill_cp_mode="in-seq-split", - disaggregation_mode="null", + disaggregation_mode=None, page_size=64, enable_hisparse=False, hicache_storage_backend="mooncake",