[mem_cache][6/N] refactor: move MHA host-pool into pool_host/mha.py (#30249)

python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py

@@ -18,10 +18,8 @@ from sglang.srt.mem_cache.memory_pool import (
     MLATokenToKVPool,
     ReqToTokenPool,
 )
-from sglang.srt.mem_cache.memory_pool_host import (
-    MLATokenToKVPoolHost,
-    get_mha_host_pool_cls,
-)
+from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost
+from sglang.srt.mem_cache.pool_host.mha import get_mha_host_pool_cls
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils.common import ceil_align
 

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -47,10 +47,8 @@ from sglang.srt.mem_cache.memory_pool import (
     MiniMaxSparseKVPool,
     MLATokenToKVPool,
 )
-from sglang.srt.mem_cache.memory_pool_host import (
-    MLATokenToKVPoolHost,
-    get_mha_host_pool_cls,
-)
+from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost
+from sglang.srt.mem_cache.pool_host.mha import get_mha_host_pool_cls
 from sglang.srt.mem_cache.radix_cache import (
     RadixCache,
     RadixKey,

python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py

@@ -19,9 +19,11 @@ from sglang.srt.mem_cache.memory_pool_host import (
     HostPoolGroup,
     LogicalHostPool,
     MambaPoolHost,
-    MHATokenToKOnlyPoolHost,
     MLATokenToKVPoolHost,
     PoolEntry,
+)
+from sglang.srt.mem_cache.pool_host.mha import (
+    MHATokenToKOnlyPoolHost,
     get_mha_host_pool_cls,
 )
 from sglang.srt.mem_cache.unified_cache_components import ComponentType

python/sglang/srt/mem_cache/kv_cache_builder.py

@@ -89,10 +89,8 @@ def maybe_register_hicache_draft(
         MHATokenToKVPool,
         MLATokenToKVPool,
     )
-    from sglang.srt.mem_cache.memory_pool_host import (
-        MLATokenToKVPoolHost,
-        get_mha_host_pool_cls,
-    )
+    from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost
+    from sglang.srt.mem_cache.pool_host.mha import get_mha_host_pool_cls
 
     pool = draft_kv_pool
     if isinstance(pool, HybridLinearKVPool):

python/sglang/srt/mem_cache/pool_host/base.py

@@ -24,6 +24,8 @@ _is_hip = is_hip()
 # Host RAM to leave free when sizing HiCache pools (OS, other processes).
 HICACHE_HOST_MEMORY_RESERVE_BYTES: int = 10 * (1024**3)
 
+_WRITE_BACK_STAGING_PAGE_CHUNK = 64
+
 
 def sync_fixed_hicache_size(size: int, host_size: int) -> int:
     """Sync fixed-size HiCache token capacity across PP ranks.

python/sglang/srt/mem_cache/storage/aibrix_kvcache/unit_test.py

@@ -8,7 +8,7 @@ from aibrix_kvcache_storage import AibrixKVCacheStorage
 
 from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig
 from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
-from sglang.srt.mem_cache.memory_pool_host import MHATokenToKVPoolHost
+from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost
 
 logging.basicConfig(
     level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"