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

fix: implement CP HiCache write retry eviction

Browse Source
This commit is contained in:
谢骞
2026-05-08 01:35:55 +08:00
parent 068aa2f910
commit 08a95a35a6
2 changed files with 94 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1073,6 +1073,56 @@ class HiRadixCache(RadixCache):
self._delete_leaf(node)
return num_evicted
def _remove_host_leaf(self, node: TreeNode) -> TreeNode:
parent = node.parent
key = self.get_child_key_fn(node.key)
v = parent.children.pop(key, None)
assert v == node, f"parent does not have child key, {key}"
if node in self.evictable_host_leaves:
self.evictable_host_leaves.remove(node)
self._update_host_leaf_status(parent)
return parent
def _evict_host_for_physical_slots(self, required_host_slots: int) -> int:
leaves = list(self.evictable_host_leaves)
eviction_heap = [
(self.eviction_strategy.get_priority(node), node) for node in leaves
]
heapq.heapify(eviction_heap)
num_evicted = 0
while num_evicted < required_host_slots and len(eviction_heap):
_priority, x = heapq.heappop(eviction_heap)
if x == self.root_node:
break
if not x.evicted:
continue
if x.pin_expiry > 0 and time.monotonic() > x.pin_expiry:
self._clear_pin(x)
if x.host_ref_counter > 0:
continue
host_indices = (
x.cp_hicache.host_indices if x.cp_hicache is not None else None
)
physical_count = len(host_indices) if host_indices is not None else 0
self._record_remove_event(x)
if physical_count > 0:
num_evicted += self.cache_controller.evict_host(host_indices)
x.host_len = 0
x.cp_hicache = None
x.host_value = None
parent = self._remove_host_leaf(x)
if len(parent.children) == 0 and parent.evicted:
new_priority = self.eviction_strategy.get_priority(parent)
heapq.heappush(eviction_heap, (new_priority, parent))
return num_evicted
def evict_host(self, num_tokens: int):
leaves = list(self.evictable_host_leaves)
eviction_heap = [
@@ -1102,16 +1152,11 @@ class HiRadixCache(RadixCache):
self._record_remove_event(x)
num_evicted += self.cache_controller.evict_host(x.host_value)
key = self.get_child_key_fn(x.key)
v = x.parent.children.pop(key, None)
assert v == x, f"parent does not have child key, {key}"
if x in self.evictable_host_leaves:
self.evictable_host_leaves.remove(x)
self._update_host_leaf_status(x.parent)
parent = self._remove_host_leaf(x)
if len(x.parent.children) == 0 and x.parent.evicted:
new_priority = self.eviction_strategy.get_priority(x.parent)
heapq.heappush(eviction_heap, (new_priority, x.parent))
if len(parent.children) == 0 and parent.evicted:
new_priority = self.eviction_strategy.get_priority(parent)
heapq.heappush(eviction_heap, (new_priority, parent))
def load_back(
self, node: TreeNode, mem_quota: Optional[int] = None

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

@@ -11,7 +11,7 @@ for _mod in ("sgl_kernel", "sgl_kernel.kvcacheio"):
sys.modules[_mod] = MagicMock()
from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata, HiRadixCache
from sglang.srt.mem_cache.radix_cache import TreeNode
from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -159,6 +159,7 @@ class FakeWriteController:
self.required_host_slots = required_host_slots
self.calls = 0
self.write_policy = "write_through"
self.evicted_host_indices = []
def write(self, device_indices, node_id=-1, priority=None):
self.calls += 1
@@ -172,6 +173,15 @@ class FakeWriteController:
)
)
def evict_host(self, host_indices):
self.evicted_host_indices.append(host_indices.clone())
return len(host_indices)
class FakeEvictionStrategy:
def get_priority(self, node):
return 0
class TestHiRadixCacheCPBackup(CustomTestCase):
def test_node_backuped_uses_cp_metadata(self):
@@ -212,18 +222,42 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
def test_write_backup_retries_by_required_physical_slots(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = FakeWriteController(required_host_slots=3)
cache._evict_host_for_physical_slots = lambda required: setattr(
cache, "evicted_required", required
)
cache.cache_controller = FakeWriteController(required_host_slots=1)
cache.evictable_host_leaves = set()
cache.eviction_strategy = FakeEvictionStrategy()
cache.get_child_key_fn = lambda key: key.token_ids[0]
cache._record_remove_event = lambda node: None
cache.ongoing_write_through = {}
cache.inc_node_lock_ref = lambda node: None
root = TreeNode()
root.key = RadixKey(token_ids=[], extra_key=None)
root.value = []
cache.root_node = root
evictable_node = TreeNode()
evictable_node.parent = root
evictable_node.key = RadixKey(token_ids=[1], extra_key=None)
evictable_node.value = None
evictable_node.host_len = 4
evictable_node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=4,
owned_positions=torch.tensor([0], dtype=torch.int64),
host_indices=torch.tensor([55], dtype=torch.int64),
)
root.children[1] = evictable_node
cache.evictable_host_leaves.add(evictable_node)
node = TreeNode()
node.value = torch.arange(16, dtype=torch.int64)
cache.write_backup(node)
self.assertEqual(cache.evicted_required, 3)
self.assertEqual(
cache.cache_controller.evicted_host_indices[0].tolist(), [55]
)
self.assertEqual(evictable_node.host_len, 0)
self.assertIsNone(evictable_node.cp_hicache)
self.assertNotIn(1, root.children)
self.assertEqual(node.host_len, 16)