From 4ea17169b0d06efb98932b5bbd453ec44b40b6a4 Mon Sep 17 00:00:00 2001 From: Yiqi Yang <134501931+decajoin@users.noreply.github.com> Date: Mon, 27 Jul 2026 13:53:02 +0800 Subject: [PATCH] [UnifiedTree] fix: drop prefetched host refill under an un-backed-up parent (#31902) Co-authored-by: hzh0425 --- .../sglang/srt/mem_cache/base_prefix_cache.py | 1 + .../unified_cache/unified_tree_core.py | 10 + .../srt/mem_cache/unified_radix_cache.py | 54 ++-- .../test_unified_radix_cache_unittest.py | 271 ++++++++++++++++-- 4 files changed, 297 insertions(+), 39 deletions(-) diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 244eef333..7a9a2b8c8 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -83,6 +83,7 @@ class InsertResult: last_device_node: Any = None mamba_exist: bool = False inserted_host_node: Any = None + host_insert_dropped: bool = False # Controller-applied actions from the non-stepped channels (e.g. insert_host); the stepped insert emits via InsertStepResult.actions. cache_actions: list[CacheAction | ComponentAction] = dataclasses.field( default_factory=list diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index df086e04c..6fe91f08e 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -1494,6 +1494,16 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): result.inserted_host_node = node.id return result + # Drop the refill only under write-through (a non-write-back policy). + if node is not self.root_node and not node.backuped and not self.is_write_back: + logger.info( + "HiCache prefetch dropped %d-token refill under un-backed-up node %d", + len(host_value), + node.id, + ) + result.host_insert_dropped = True + return result + new_node = self._new_node(priority=node.priority) new_node.parent = node new_node.key = key diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 69477a27a..09d58daad 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -1281,39 +1281,51 @@ class UnifiedRadixCache(BasePrefixCache): # Apply the host-insert walk's actions before the transfer commit. self._apply_cache_actions(insert_result.cache_actions) - commit_actions: list[CacheAction | ComponentAction] = [] - self.tree_core.commit_hicache_transfers( - last_host_node_id, - CacheTransferPhase.PREFETCH, - comp_xfers, - cache_actions=commit_actions, - insert_result=insert_result, - pool_storage_result=operation.pool_storage_result, - ) - self._apply_cache_actions(commit_actions) - # The commit emits via commit_actions only; the walk's were applied above. - assert not insert_result.cache_actions - self.cache_controller.mem_pool_host.free( - host_indices[: insert_result.prefix_len] - ) - self.cache_controller.append_host_mem_release( - host_indices[min_completed_tokens:completed_tokens] - ) + if insert_result.host_insert_dropped: + self.cache_controller.append_host_mem_release( + host_indices=host_indices[:completed_tokens], + extra_pools=[x for xfers in comp_xfers.values() for x in xfers], + ) + loaded_from_storage = 0 + released_tokens = completed_tokens + else: + commit_actions: list[CacheAction | ComponentAction] = [] + self.tree_core.commit_hicache_transfers( + last_host_node_id, + CacheTransferPhase.PREFETCH, + comp_xfers, + cache_actions=commit_actions, + insert_result=insert_result, + pool_storage_result=operation.pool_storage_result, + ) + self._apply_cache_actions(commit_actions) + # The commit emits via commit_actions only; the walk's were applied above. + assert not insert_result.cache_actions + + self.cache_controller.mem_pool_host.free( + host_indices[: insert_result.prefix_len] + ) + self.cache_controller.append_host_mem_release( + host_indices[min_completed_tokens:completed_tokens] + ) + loaded_from_storage = min_completed_tokens - insert_result.prefix_len + released_tokens = completed_tokens - min_completed_tokens + self.dec_host_lock_ref(last_host_node_id, anchor_lock_params) del self.ongoing_prefetch[req_id] self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key) - loaded_from_storage = min_completed_tokens - insert_result.prefix_len self.prefetch_loaded_tokens_by_reqid[req_id] = loaded_from_storage logger.info( - "HiCache prefetch success req=%s completed_local=%d completed_synced=%d matched=%d loaded=%d tail_release=%d occupied=%d", + "HiCache prefetch %s req=%s completed_local=%d completed_synced=%d matched=%d loaded=%d released=%d occupied=%d", + "dropped" if insert_result.host_insert_dropped else "success", req_id, completed_tokens, min_completed_tokens, insert_result.prefix_len, loaded_from_storage, - completed_tokens - min_completed_tokens, + released_tokens, self.cache_controller.prefetch_tokens_occupied, ) if self.enable_storage_metrics and self.storage_metrics_collector is not None: diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 0830ea6d9..c06e34295 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -37,7 +37,11 @@ from sglang.srt.mem_cache.base_prefix_cache import ( ) from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.common import available_and_evictable_str -from sglang.srt.mem_cache.hicache_storage import PoolName +from sglang.srt.mem_cache.hicache_storage import ( + PoolName, + PoolTransfer, + PoolTransferResult, +) from sglang.srt.mem_cache.memory_pool import ( HybridLinearKVPool, HybridReqToTokenPool, @@ -74,6 +78,7 @@ from sglang.srt.mem_cache.unified_radix_cache import ( UnifiedLRUList, UnifiedRadixCache, UnifiedTreeNode, + _OngoingPrefetch, _OngoingWriteThrough, ) from sglang.srt.runtime_context import get_server_args @@ -3658,9 +3663,10 @@ class UnifiedRadixCacheSuite: sw = cache.sliding_window_size swa = cache.components[ComponentType.SWA] # pool can't satisfy even after evict -> participates but aborts (no buffer) - with mock.patch.object( - cache.swa_kv_pool_host, "alloc", return_value=None - ), mock.patch.object(cache, "evict_host", autospec=True) as evict_host: + with ( + mock.patch.object(cache.swa_kv_pool_host, "alloc", return_value=None), + mock.patch.object(cache, "evict_host", autospec=True) as evict_host, + ): self.assertTrue( swa.prepare_prefetch( cache.root_node.id, prefetch_tokens=sw @@ -4315,9 +4321,14 @@ class UnifiedRadixCacheSuite: retry_slot = req_to_token_pool.mamba_allocator.alloc(1) # first alloc fails -> prepare must evict a mamba slot and retry - with mock.patch.object( - req_to_token_pool.mamba_allocator, "alloc", side_effect=[None, retry_slot] - ), mock.patch.object(cache, "evict", autospec=True) as evict: + with ( + mock.patch.object( + req_to_token_pool.mamba_allocator, + "alloc", + side_effect=[None, retry_slot], + ), + mock.patch.object(cache, "evict", autospec=True) as evict, + ): prep = comp.prepare_load_back(leaf.id, req=req) evict.assert_called_once_with(EvictParams(num_tokens=0, mamba_num=1)) self.assertIs(prep.allocated_mamba_slot, retry_slot) @@ -5992,9 +6003,12 @@ class TestReturnedValuesDrain(_InsertWalkSuite): device_frees.clear() host_frees.clear() - with mock.patch.object( - cache.tree_core, name, return_value=make_result() - ), mock.patch.object(cache, "_free_values", side_effect=record): + with ( + mock.patch.object( + cache.tree_core, name, return_value=make_result() + ), + mock.patch.object(cache, "_free_values", side_effect=record), + ): returned = call() self.assertEqual(returned, expected) ((device_frees, host_frees),) = drained @@ -6022,11 +6036,16 @@ class TestReturnedValuesDrain(_InsertWalkSuite): ) for ct in order } - with mock.patch.object( - cache, - "_apply_cache_action", - side_effect=lambda action: freed.append(("device", action.component_type)), - ), mock.patch.dict(cache.components, fake_components): + with ( + mock.patch.object( + cache, + "_apply_cache_action", + side_effect=lambda action: freed.append( + ("device", action.component_type) + ), + ), + mock.patch.dict(cache.components, fake_components), + ): cache._free_values(device_frees, host_frees) self.assertEqual( @@ -6051,9 +6070,10 @@ class TestReturnedValuesDrain(_InsertWalkSuite): raise RuntimeError("boom") host_mock = mock.MagicMock() - with mock.patch.object( - cache, "_apply_cache_action", side_effect=boom_on_swa - ), mock.patch.dict(cache.components, {ComponentType.FULL: host_mock}): + with ( + mock.patch.object(cache, "_apply_cache_action", side_effect=boom_on_swa), + mock.patch.dict(cache.components, {ComponentType.FULL: host_mock}), + ): with self.assertRaises(RuntimeError): cache._free_values(device_frees, host_frees) @@ -6138,6 +6158,7 @@ class TestPrefetchCommitOrdering(CustomTestCase): insert_result = mock.MagicMock() insert_result.cache_actions = [walk_action] insert_result.prefix_len = 4 + insert_result.host_insert_dropped = False cache.tree_core.insert_host.return_value = insert_result cache.ongoing_prefetch = { "req": ( @@ -6178,5 +6199,219 @@ class TestPrefetchCommitOrdering(CustomTestCase): self.assertEqual(cache.ongoing_prefetch, {}) +class TestUnifiedRadixPrefetchCorruption(CustomTestCase): + """Prefetch must not hang a backed-up host child under an un-backed-up parent. + + Under write-through that broke the "child backed up => parent backed up" + invariant, failing as an idle-sanity error and, on eviction, as + `_remove_leaf_from_parent -> assert v == node`. Fix: drop the refill. + """ + + ps = 16 + cfg = CacheConfig( + page_size=ps, + components=(ComponentType.FULL,), + kv_size=4096, + max_context_len=4096, + ) + + def _init_hicache(self, cache, *, write_policy="write_through"): + import sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler as assembler + + orig_get_mha_host_pool_cls = assembler.get_mha_host_pool_cls + + def get_mha_host_pool_cls_wrapper(device_pool): + host_pool_cls = orig_get_mha_host_pool_cls(device_pool) + + def kv_host_pool_wrapper(*args, **kwargs): + kwargs["pin_memory"] = False + return host_pool_cls(*args, **kwargs) + + return kv_host_pool_wrapper + + patcher = mock.patch.object( + assembler, + "get_mha_host_pool_cls", + side_effect=get_mha_host_pool_cls_wrapper, + ) + patcher.start() + self.addCleanup(patcher.stop) + + server_args = ServerArgs( + model_path="dummy", + page_size=self.cfg.page_size, + hicache_io_backend="direct", + hicache_write_policy=write_policy, + ) + server_args._mamba_cache_chunk_size = max(FLA_CHUNK_SIZE, self.cfg.page_size) + set_global_server_args_for_scheduler(server_args) + cache.init_hicache(server_args, cache.cache_init_params) + cache.write_through_threshold = 1 << 30 + cache.load_back_threshold = 0 + + def _insert_device(self, cache, allocator, ids): + """Insert a device-only chain (no auto-backup) and return its leaf id.""" + key = RadixKey(array("q", ids)).page_aligned(self.ps) + val = allocator.alloc(len(key)) + self.assertIsNotNone(val) + val = val.to(dtype=torch.int64) + cache.insert(InsertParams(key=key, value=val, prev_prefix_len=0)) + return cache.match_prefix(MatchPrefixParams(key=key)).last_device_node + + def _attach_host_child(self, cache, parent_id, start_token): + """Mimic a prefetch commit: hang a backed-up host chain under `parent_id`.""" + ps = self.ps + child_key = RadixKey( + array("q", list(range(start_token, start_token + 2 * ps))) + ).page_aligned(ps) + host_idx = cache.cache_controller.mem_pool_host.alloc(len(child_key)) + self.assertIsNotNone(host_idx, "host pool alloc failed") + host_idx = host_idx.to(dtype=torch.int64) + hashes = [f"h{i}" for i in range(len(child_key) // ps)] + res = cache.tree_core.insert_host(parent_id, child_key, host_idx, hashes) + if res.host_insert_dropped: + cache.cache_controller.mem_pool_host.free(host_idx) + return res.inserted_host_node + + def test_prefetch_refill_under_unbacked_parent_is_dropped(self): + """Write-through: a refill under an un-backed-up parent is dropped.""" + cache, allocator, _ = build_fixture(self.cfg) + self._init_hicache(cache) + + parent_id = self._insert_device( + cache, allocator, list(range(1, 1 + 3 * self.ps)) + ) + parent = cache.tree_core.node_by_id(parent_id) + self.assertFalse(parent.backuped) + + child = self._attach_host_child(cache, parent_id, start_token=1000) + self.assertIsNone(child) + self.assertEqual(len(parent.children), 0) + cache.sanity_check() + + def test_dropped_prefetch_releases_all_host_resources(self): + """The caller owns every completed buffer when host insertion drops.""" + cache, allocator, _ = build_fixture(self.cfg) + self._init_hicache(cache) + + parent_id = self._insert_device( + cache, allocator, list(range(1, 1 + 3 * self.ps)) + ) + prefetch_key = RadixKey( + array("q", list(range(1000, 1000 + 2 * self.ps))) + ).page_aligned(self.ps) + completed_tokens = len(prefetch_key) + host_indices = cache.cache_controller.mem_pool_host.alloc(completed_tokens) + self.assertIsNotNone(host_indices) + + swa_transfer = PoolTransfer( + name=PoolName.SWA, + host_indices=torch.arange(self.ps, dtype=torch.int64), + ) + mamba_transfer = PoolTransfer( + name=PoolName.MAMBA, + host_indices=torch.arange(1, dtype=torch.int64), + ) + swa_component = mock.Mock() + mamba_component = mock.Mock() + cache.tree_core.components_by_type[ComponentType.SWA] = swa_component + cache.tree_core.components_by_type[ComponentType.MAMBA] = mamba_component + comp_xfers = { + ComponentType.SWA: [swa_transfer], + ComponentType.MAMBA: [mamba_transfer], + } + + operation = mock.Mock() + operation.host_indices = host_indices + operation.pool_storage_result = PoolTransferResult( + kv_hit_pages=completed_tokens // self.ps, + extra_pool_hit_pages={ + PoolName.SWA: 1, + PoolName.MAMBA: 1, + }, + ) + anchor_lock_params = cache.inc_host_lock_ref(parent_id).to_dec_params() + req_id = "drop-all-resources" + cache.ongoing_prefetch[req_id] = _OngoingPrefetch( + parent_id, + prefetch_key, + host_indices, + operation, + anchor_lock_params, + comp_xfers, + ) + cache.cache_controller.prefetch_tokens_occupied = completed_tokens + hashes = [f"h{i}" for i in range(completed_tokens // self.ps)] + + with ( + mock.patch.object(cache, "can_terminate_prefetch", return_value=True), + # Isolate the drop-release branch under test from the hybrid-sync + # step: treat the whole fetched prefix as usable so the insert runs. + mock.patch.object( + cache, + "_sync_and_check_hybrid_prefetch_result", + return_value=completed_tokens, + ), + mock.patch.object( + cache.cache_controller, + "terminate_prefetch", + return_value=(completed_tokens, hashes), + ), + # No storage backend in this fixture, so the real release queues + # don't exist; assert on the release call instead of draining them. + mock.patch.object( + cache.cache_controller, "append_host_mem_release" + ) as release, + ): + self.assertTrue(cache.check_prefetch_progress(req_id)) + + swa_component.commit_hicache_transfer.assert_not_called() + mamba_component.commit_hicache_transfer.assert_not_called() + self.assertEqual(cache.pop_prefetch_loaded_tokens(req_id), 0) + self.assertEqual(len(cache.tree_core.node_by_id(parent_id).children), 0) + + drop_releases = [ + call + for call in release.call_args_list + if call.kwargs.get("extra_pools") is not None + ] + self.assertEqual(len(drop_releases), 1) + self.assertTrue( + torch.equal( + drop_releases[0].kwargs["host_indices"], + host_indices[:completed_tokens], + ) + ) + self.assertIs(drop_releases[0].kwargs["extra_pools"][0], swa_transfer) + self.assertIs(drop_releases[0].kwargs["extra_pools"][1], mamba_transfer) + + cache.sanity_check() + + def test_prefetch_refill_leaves_eviction_path_uncorrupted(self): + """Write-through: eviction after such a refill must not corrupt the tree.""" + cache, allocator, _ = build_fixture(self.cfg) + self._init_hicache(cache) + + parent_id = self._insert_device( + cache, allocator, list(range(1, 1 + 3 * self.ps)) + ) + self._attach_host_child(cache, parent_id, start_token=1000) + + cache.evict(EvictParams(num_tokens=10 * self.ps)) + cache.sanity_check() + + def test_prefetch_refill_kept_under_unbacked_parent_in_write_back(self): + """Write-back keeps the refill (it has no backed-up-parent requirement).""" + cache, allocator, _ = build_fixture(self.cfg) + self._init_hicache(cache, write_policy="write_back") + + parent_id = self._insert_device( + cache, allocator, list(range(1, 1 + 3 * self.ps)) + ) + child = self._attach_host_child(cache, parent_id, start_token=1000) + self.assertIsNotNone(child) + cache.sanity_check() + + if __name__ == "__main__": unittest.main()