[HiCache] Support host memory size smaller than device memory size

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

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]:

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()

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)

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)