chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

fix: address CP HiCache review regressions

Browse Source
This commit is contained in:
谢骞
2026-05-08 04:11:25 +08:00
parent 8844c303e2
commit f38937723e
4 changed files with 104 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
python/sglang/srt/mem_cache/hiradix_cache.py
View File

@@ -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,

6
python/sglang/srt/server_args.py
View File

@@ -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

86
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
View File

@@ -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

2
test/registered/unit/server_args/test_server_args.py
View File

@@ -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",