Refactor HiCache host pool management (#36232)

test/registered/unit/mem_cache/test_decode_retraction_backup.py

@@ -117,7 +117,6 @@ class TestDecodeRetractionBackup(unittest.TestCase):
                 device_pools=(draft_pool,),
             ),
             server_args=server_args,
-            page_size=1,
         )
         self.assertIn(PoolName.DRAFT, cache.host_pool_group.entry_map)
         cache.validate_retraction_host_capacity()

test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py

@@ -21,10 +21,9 @@ from sglang.srt.mem_cache.l2_transfer import L2Transfer, L2TransferEngine
 from sglang.srt.mem_cache.memory_pool_host import (
     DeepSeekV4PagedHostPool,
     DeepSeekV4StateHostPool,
-    HostPoolGroup,
     LogicalHostPool,
-    PoolEntry,
 )
+from sglang.srt.mem_cache.pool_host import HostPoolGroup, PoolEntry
 from sglang.srt.mem_cache.pool_host.dsa import DSAIndexerPoolHost
 from sglang.srt.mem_cache.pool_host.mamba import MambaPoolHost
 from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost
@@ -222,8 +221,6 @@ class TestHiCacheStagedWriteBackDispatch(CustomTestCase):
             op.pool_transfers,
         )
         controller.mem_pool_host = _host_group_stub([], can_use_write_back_jit=False)
-        controller.has_draft = False
-        controller.has_mtp_draft = False
         controller._l2_transfers.side_effect = lambda *args: (
             HybridCacheController._l2_transfers(controller, *args)
         )
@@ -287,20 +284,19 @@ class TestHiCacheStagedWriteBackDispatch(CustomTestCase):
     def test_packed_draft_load_is_flattened_into_l2_transfers(self):
         host_pool = mock.Mock()
         controller = HybridCacheController.__new__(HybridCacheController)
+        entry = PoolEntry(
+            name=PoolName.KV,
+            host_pool=host_pool,
+            device_pool=mock.sentinel.target_device_pool,
+            layer_mapper={0: 0, 1: 1, 2: 2}.get,
+            is_primary_index_anchor=True,
+            packed_draft_device_pools=(mock.sentinel.draft_device_pool,),
+        )
         controller.mem_pool_host = SimpleNamespace(
-            anchor_entry=PoolEntry(
-                name=PoolName.KV,
-                host_pool=host_pool,
-                device_pool=mock.sentinel.target_device_pool,
-                layer_mapper={0: 0, 1: 1, 2: 2}.get,
-                is_primary_index_anchor=True,
-            ),
-            entry_map={},
+            anchor_entry=entry,
+            entry_map={entry.name: entry},
         )
         controller.layer_num = 2
-        controller.has_mtp_draft = True
-        controller.mtp_draft_device_pools = (mock.sentinel.draft_device_pool,)
-        controller.has_draft = False
 
         self.assertEqual(
             len(controller._l2_transfers(_indices(0, 2), _indices(2, 4))), 1
@@ -937,7 +933,6 @@ class TestHiCacheStagedWriteBackDispatch(CustomTestCase):
             captured, can_use_write_back_jit=True
         )
         controller.mem_pool_device = None
-        controller.has_draft = False
         controller.ack_write_queue = []
         controller.move_hybrid_indices = mock.Mock(
             side_effect=AssertionError(
@@ -972,7 +967,6 @@ class TestHiCacheStagedWriteBackDispatch(CustomTestCase):
             captured, can_use_write_back_jit=False
         )
         controller.mem_pool_device = None
-        controller.has_draft = False
         controller.ack_write_queue = []
         controller.move_hybrid_indices = mock.Mock(
             return_value=(op.host_indices, op.device_indices, op.pool_transfers)
@@ -1007,7 +1001,6 @@ class TestHiCacheStagedWriteBackDispatch(CustomTestCase):
         controller.io_backend = "kernel"
         controller.mem_pool_host = FakeHostPool()
         controller.mem_pool_device = None
-        controller.has_draft = False
         controller.device = "cuda"
         controller.ack_write_queue = []
         controller.move_indices = mock.Mock(
@@ -1044,7 +1037,6 @@ class TestHiCacheStagedWriteBackDispatch(CustomTestCase):
         controller.io_backend = "kernel"
         controller.mem_pool_host = FakeHostPool()
         controller.mem_pool_device = None
-        controller.has_draft = False
         controller.device = "cuda"
         controller.ack_write_queue = []
         controller.move_indices = mock.Mock(

test/registered/unit/mem_cache/test_mem_pool_host.py

@@ -6,12 +6,13 @@ import unittest.mock
 
 import torch
 
+from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer
 from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
 from sglang.srt.mem_cache.memory_pool_host import (
     DeepSeekV4PagedHostPool,
     LogicalHostPool,
 )
-from sglang.srt.mem_cache.pool_host import base
+from sglang.srt.mem_cache.pool_host import HostPoolGroup, PoolEntry, base
 from sglang.srt.mem_cache.pool_host.mamba import MambaPoolHost
 from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost
 from sglang.srt.runtime_context import get_context
@@ -238,5 +239,54 @@ class TestHostMemoryBudget(CustomTestCase):
             self.assertEqual(base.ranks_per_host(), 8)
 
 
+class TestHostPoolGroup(CustomTestCase):
+    @staticmethod
+    def _group(**sizes):
+        return HostPoolGroup(
+            [
+                PoolEntry(
+                    name=PoolName(name),
+                    host_pool=LogicalHostPool(size=size, page_size=1),
+                    device_pool=None,
+                    layer_mapper=lambda layer_id: layer_id,
+                    is_primary_index_anchor=name == PoolName.KV.value,
+                )
+                for name, size in sizes.items()
+            ]
+        )
+
+    def test_resolve_and_release_multi_pool_allocation(self):
+        group = self._group(kv=4, swa=2)
+        primary = group.alloc(2)
+        transfers = [
+            PoolTransfer(name=PoolName.SWA, device_indices=torch.arange(2)),
+            PoolTransfer(name=PoolName.INDEXER, indices_from_pool=PoolName.SWA),
+        ]
+
+        self.assertIsNotNone(
+            group.resolve_host_transfers(
+                transfers,
+                primary_device_indices=torch.arange(2),
+                primary_host_indices=primary,
+            )
+        )
+        self.assertIs(transfers[1].host_indices, transfers[0].host_indices)
+        group.free(primary)
+        group.release_transfers(transfers)
+        self.assertEqual(group.available_size(), 4)
+        self.assertEqual(group.available_size(PoolName.SWA), 2)
+
+    def test_resolve_rolls_back_partial_allocation(self):
+        group = self._group(kv=4, swa=2, mamba=1)
+        transfers = [
+            PoolTransfer(name=PoolName.SWA, device_indices=torch.arange(2)),
+            PoolTransfer(name=PoolName.MAMBA, device_indices=torch.arange(2)),
+        ]
+
+        self.assertIsNone(group.resolve_host_transfers(transfers))
+        self.assertIsNone(transfers[0].host_indices)
+        self.assertEqual(group.available_size(PoolName.SWA), 2)
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -5518,7 +5518,7 @@ class UnifiedRadixCacheSuite:
         self.assertEqual(xfer.nodes_to_load, [n.id for n in loaded_nodes])
 
         # Allocate SWA device slots from the inner allocator (mirrors how
-        # _resolve_pool_transfers_allocation routes via device_alloc_fn ->
+        # _resolve_device_transfers routes via device_alloc_fn ->
         # swa_attn_allocator.alloc on the load-back path).
         n_swa = int(xfer.host_indices.numel())
         new_swa = allocator.swa_attn_allocator.alloc(n_swa)

test/registered/unit/test_supplied_instance_exposure_ratchet.py

@@ -219,7 +219,6 @@ _OVERRIDDEN_AND_READ = {
     ("weight_cache/daemon.py", "model_path"),
     ("configs/model_config.py", "dtype"),
     ("configs/model_config.py", "model_path"),
-    ("mem_cache/kv_cache_builder.py", "hicache_storage_backend"),
     ("mem_cache/pool_host/common.py", "hicache_storage_backend"),
     ("mem_cache/pool_host/common.py", "hicache_storage_backend_extra_config"),
     ("mem_cache/unified_radix_cache.py", "hicache_storage_backend"),