From 3d9f693307574e1307cee8bbf0a57cbfbbea5758 Mon Sep 17 00:00:00 2001 From: Teng Ma Date: Mon, 12 Jan 2026 00:09:48 +0800 Subject: [PATCH 1/4] [HiCache] Support host memory size smaller than device memory size --- .../sglang/srt/managers/cache_controller.py | 10 +++++-- python/sglang/srt/mem_cache/hiradix_cache.py | 30 +++++++++++-------- .../sglang/srt/mem_cache/memory_pool_host.py | 6 ++-- 3 files changed, 28 insertions(+), 18 deletions(-) diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 0caa63cde0e2..09027dd2d247 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -302,9 +302,13 @@ def __init__( self.enable_storage = True # todo: threshold policy for prefetching self.prefetch_threshold = max(prefetch_threshold, self.page_size) - self.prefetch_capacity_limit = int( - 0.8 * (self.mem_pool_host.size - self.mem_pool_device.size) - ) + if self.mem_pool_host.size > self.mem_pool_device.size: + self.prefetch_capacity_limit = int( + 0.8 * (self.mem_pool_host.size - self.mem_pool_device.size) + ) + else: + self.prefetch_capacity_limit = int(0.5 * self.mem_pool_host.size) + # granularity of batch storage IO operations, in number of pages self.storage_batch_size = 128 # tracking the number of tokens locked in prefetching, updated by the main scheduler thread diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index f6cfca8b6fad..b956e66aa23e 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -345,8 +345,12 @@ def evict(self, num_tokens: int): if not x.backuped: if self.cache_controller.write_policy == "write_back": # write to host if the node is not backuped - num_evicted += self.write_backup(x, write_back=True) - write_back_nodes.append(x) + backed_up_len = self.write_backup(x, write_back=True) + if backed_up_len > 0: + num_evicted += backed_up_len + write_back_nodes.append(x) + else: + num_evicted += self._evict_regular(x) else: num_evicted += self._evict_regular(x) else: @@ -397,23 +401,25 @@ def evict_host(self, num_tokens: int): _priority, x = heapq.heappop(eviction_heap) if x == self.root_node: break - # only evict the host value of evicted nodes - if not x.evicted: - continue - + # node is protected from eviction as it has ongoing prefetch or backup to storage if x.host_ref_counter > 0: continue + if x.host_value is None: + continue + num_evicted += self.cache_controller.evict_host(x.host_value) + x.host_value = None - 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.evicted: + 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 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(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)) def load_back( self, node: TreeNode, mem_quota: Optional[int] = None diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 548c9d9f18cb..99a4c21dd596 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -158,9 +158,9 @@ def __init__( self.start_layer = device_pool.start_layer self.end_layer = device_pool.end_layer - assert ( - self.size > device_pool.size - ), "The host memory should be larger than the device memory with the current protocol" + # assert ( + # self.size > device_pool.size + # ), "The host memory should be larger than the device memory with the current protocol" # Verify there is enough available host memory. host_mem = psutil.virtual_memory() From cbfb249e7408d84540e141643a6acebbb044dbf7 Mon Sep 17 00:00:00 2001 From: Teng Ma Date: Mon, 12 Jan 2026 01:07:39 +0800 Subject: [PATCH 2/4] Update python/sglang/srt/mem_cache/hiradix_cache.py --- python/sglang/srt/mem_cache/hiradix_cache.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index b956e66aa23e..337571a293b0 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -401,7 +401,6 @@ def evict_host(self, num_tokens: int): _priority, x = heapq.heappop(eviction_heap) if x == self.root_node: break - # node is protected from eviction as it has ongoing prefetch or backup to storage if x.host_ref_counter > 0: continue From 35fd328128d2f657d5ccb47d06ba2cb811aab146 Mon Sep 17 00:00:00 2001 From: Teng Ma Date: Wed, 25 Feb 2026 16:43:27 +0800 Subject: [PATCH 3/4] update --- .../hicache/test_hicache_storage_file_backend.py | 9 +++++++++ .../test_hicache_storage_mooncake_backend.py | 13 +++++++++++++ 2 files changed, 22 insertions(+) diff --git a/test/registered/hicache/test_hicache_storage_file_backend.py b/test/registered/hicache/test_hicache_storage_file_backend.py index ed9fca934046..72aca24ed148 100644 --- a/test/registered/hicache/test_hicache_storage_file_backend.py +++ b/test/registered/hicache/test_hicache_storage_file_backend.py @@ -330,5 +330,14 @@ def run_eval_accuracy_test(test_instance, accuracy_threshold: float = 0.03): ) +class TestHiCacheStorageSmallHostMemory(HiCacheStorageBaseMixin, CustomTestCase): + """Test HiCache file backend when host (L2) memory is smaller than device (L1) memory""" + + @classmethod + def _get_additional_server_args_and_env(cls): + server_args = {"--hicache-ratio": 0.5} + return server_args, {"SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir} + + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/test/registered/hicache/test_hicache_storage_mooncake_backend.py b/test/registered/hicache/test_hicache_storage_mooncake_backend.py index 947bce776225..232b773e74c7 100644 --- a/test/registered/hicache/test_hicache_storage_mooncake_backend.py +++ b/test/registered/hicache/test_hicache_storage_mooncake_backend.py @@ -282,5 +282,18 @@ def test_eval_accuracy(self): run_eval_accuracy_test(self) +class TestMooncakeBackendSmallHostMemory( + HiCacheStorageMooncakeBackendBaseMixin, CustomTestCase +): + """Test HiCache-Mooncake backend when host (L2) memory is smaller than device (L1) memory""" + + @classmethod + def _get_additional_server_args_and_env(cls): + server_args, env_vars = super()._get_additional_server_args_and_env() + server_args["--hicache-ratio"] = 0.5 + server_args["--hicache-mem-layout"] = "page_first" + return server_args, env_vars + + if __name__ == "__main__": unittest.main(verbosity=2) From f6005cfea35cfd3636fae91bc05d48871efa7813 Mon Sep 17 00:00:00 2001 From: Xuchun Shang Date: Tue, 3 Mar 2026 12:12:54 +0800 Subject: [PATCH 4/4] fix lint Signed-off-by: Xuchun Shang --- python/sglang/srt/mem_cache/hiradix_cache.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 7cf3757f7fb6..510742ed3aee 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -980,7 +980,7 @@ def evict_host(self, num_tokens: int): # node is protected from eviction as it has ongoing prefetch, backup, or pin if x.host_ref_counter > 0: continue - + if x.host_value is None: continue