diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 51030ae1d1bd..b9bc8b95d1f2 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -448,9 +448,12 @@ def attach_storage_backend( self.enable_storage = True # todo: threshold policy for prefetching self.prefetch_threshold = max(prefetch_threshold, self.page_size) - self.prefetch_capacity_limit = max( - 0, 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 2acb01157931..510742ed3aee 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -907,8 +907,12 @@ def evict(self, params: EvictParams) -> EvictResult: 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: @@ -977,22 +981,28 @@ def evict_host(self, num_tokens: int): if x.host_ref_counter > 0: continue + if x.host_value is None: + continue + # 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) + x.host_value = None + + 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)) - 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) - 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 ) -> Optional[torch.Tensor]: diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 0b8199d9f7e9..80852e3c5961 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -165,9 +165,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() 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)