minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Write-back policy fix for unified tree (#31845)

Browse Source
This commit is contained in:
Zhiqiang Xie
2026-07-24 09:58:18 -07:00
committed by GitHub
parent 448662e85e
commit 5da0b6ec39
3 changed files with 164 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+75 -14
View File
@@ -388,6 +388,7 @@ class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase):
self.mem_pool_host = FakeHostPool()
self.prefetch_tokens_occupied = 0
self.prefetch_args = None
self.write_policy = "write_through"
def prefetch_rate_limited(self):
return False
@@ -2734,34 +2735,94 @@ class UnifiedRadixCacheSuite:
[conv[:, mamba_indices].float().cpu().clone() for conv in mamba_cache.conv],
)
def test_hicache_evict_device_leaf_aborts_demote_when_backup_fails(self):
"""when write_backup cannot allocate host pool,
_evict_device_leaf should not evict it to host."""
def test_hicache_write_back_evict_drops_unbacked_leaf_when_host_full(self):
"""Write-back eviction will keep freeing device KV when the host pool
is exhausted and host eviction cannot free space to prevent OOM."""
if self._skip_unsupported_hicache_test():
return
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
self._init_hicache(cache, write_policy="write_back")
ct = ComponentType.FULL
seq = self._make_seq(1, 2)
# Two-node chain: the drop must cascade parent-ward across eviction
# iterations, not just delete a single leaf.
seq_parent = self._make_seq(1, 2)
self._insert(cache, allocator, req_to_token_pool, seq_parent)
seq = seq_parent + self._make_seq(1000, 1)
self._insert(cache, allocator, req_to_token_pool, seq)
m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
node = m.last_device_node
self.assertIsNot(node, cache.root_node)
self.assertFalse(node.backuped)
self.assertFalse(node.evicted)
tracker = {c: 0 for c in cache.tree_components}
with mock.patch.object(cache, "write_backup", return_value=0):
cache._evict_device_leaf(node, tracker)
# Exhaust the KV host pool. The tree has no host leaves, so
# evict_host cannot free anything and every D->H backup fails.
host_pool = cache.cache_controller.mem_pool_host
self.assertIsNotNone(host_pool.alloc(host_pool.available_size()))
self.assertEqual(host_pool.available_size(), 0)
self.assertFalse(node.evicted)
self.assertIsNotNone(node.component_data[ct].value)
self.assertIsNone(node.component_data[ct].host_value)
result = cache.evict(EvictParams(num_tokens=len(seq)))
self.assertGreaterEqual(result.num_tokens_evicted, len(seq))
with self.assertRaises(AssertionError):
cache._evict_to_host(node, {c: 0 for c in cache.tree_components})
# The chain is gone entirely: no device hit, no host hit.
m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
self.assertEqual(len(m.device_indices), 0)
self.assertEqual(m.host_hit_length, 0)
cache.sanity_check()
def test_hicache_write_back_drop_respects_pins_then_frees_subtree(self):
"""The host-pressure drop fallback must decline while any node in the
unbacked subtree is host-pinned, then reclaim the whole subtree --
including a demoted child's host backup -- once unpinned."""
if self._skip_unsupported_hicache_test():
return
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
self._init_hicache(cache, write_policy="write_back")
host_pool = cache.cache_controller.mem_pool_host
baseline_host = host_pool.available_size()
parent_seq = self._make_seq(1, 2)
self._insert(cache, allocator, req_to_token_pool, parent_seq)
child_seq = parent_seq + self._make_seq(1000, 1)
self._insert(cache, allocator, req_to_token_pool, child_seq)
m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", child_seq))))
child = m.last_device_node
# Evict only the child leaf -> real backup + demote, leaving it
# host-only under a still-unbacked device parent (write-back backs up
# single nodes leaf-first, so this is a normal intermediate state).
result = cache.evict(EvictParams(num_tokens=len(child.key)))
self.assertGreaterEqual(result.num_tokens_evicted, len(child.key))
self.assertTrue(child.evicted and child.backuped)
parent = child.parent
self.assertFalse(parent.backuped)
self.assertGreater(baseline_host - host_pool.available_size(), 0)
# From here every backup fails (controller.write returns None), so
# each evict() attempts the drop fallback on the parent.
with mock.patch.object(cache.cache_controller, "write", return_value=None):
# Pinned subtree root: drop declines, chain stays intact.
cache.inc_host_lock_ref(parent)
result = cache.evict(EvictParams(num_tokens=len(parent_seq)))
self.assertEqual(result.num_tokens_evicted, 0)
cache.dec_host_lock_ref(parent)
# Pinned host-only descendant: drop declines as well.
cache.inc_host_lock_ref(child)
result = cache.evict(EvictParams(num_tokens=len(parent_seq)))
self.assertEqual(result.num_tokens_evicted, 0)
m = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", parent_seq)))
)
self.assertEqual(len(m.device_indices), len(parent_seq))
cache.dec_host_lock_ref(child)
# Unpinned: the subtree drops and the child's host slots return.
result = cache.evict(EvictParams(num_tokens=len(parent_seq)))
self.assertGreaterEqual(result.num_tokens_evicted, len(parent_seq))
self.assertEqual(host_pool.available_size(), baseline_host)
m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", child_seq))))
self.assertEqual(len(m.device_indices), 0)
self.assertEqual(m.host_hit_length, 0)
cache.sanity_check()
def test_hicache_evict_to_host(self):