Fix --hicache-size allocating ~2x host memory on hybrid SWA (#32373)

Co-authored-by: cctry <cctry@fb.com>
Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>

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

@@ -64,6 +64,7 @@ def build_kv_host_pool(
     server_args: ServerArgs,
     use_mla: bool,
     override_kv_cache_dim: Optional[int] = None,
+    host_size: Optional[float] = None,
 ):
     kv_host_pool_cls = (
         MLATokenToKVPoolHost if use_mla else get_mha_host_pool_cls(kv_pool)
@@ -74,7 +75,7 @@ def build_kv_host_pool(
     return kv_host_pool_cls(
         kv_pool,
         server_args.hicache_ratio,
-        server_args.hicache_size,
+        server_args.hicache_size if host_size is None else host_size,
         page_size,
         server_args.hicache_mem_layout,
         allocator_type=_get_allocator_type(server_args),
@@ -82,6 +83,22 @@ def build_kv_host_pool(
     )
 
 
+def _split_hicache_size(
+    hicache_size: int, kv_pools: tuple[Any, ...]
+) -> tuple[float, ...]:
+    device_pool_sizes = []
+    for kv_pool in kv_pools:
+        size_bytes = kv_pool.get_kv_size_bytes()
+        device_pool_sizes.append(
+            sum(size_bytes) if isinstance(size_bytes, tuple) else size_bytes
+        )
+    total_device_pool_size = sum(device_pool_sizes)
+    return tuple(
+        hicache_size * size_bytes / total_device_pool_size
+        for size_bytes in device_pool_sizes
+    )
+
+
 def build_pool_entry(
     *,
     name: PoolName,
@@ -182,17 +199,24 @@ def build_hybrid_swa_stack(
     enable_storage_metrics: bool = False,
 ) -> tuple[HostPoolGroup, HybridCacheController]:
     transfer_layer_num = len(full_layer_mapping | swa_layer_mapping)
+    kv_host_size = swa_host_size = None
+    if server_args.hicache_size > 0:
+        kv_host_size, swa_host_size = _split_hicache_size(
+            server_args.hicache_size, (full_kv_pool, swa_kv_pool)
+        )
     kv_host_pool = build_kv_host_pool(
         kv_pool=full_kv_pool,
         page_size=params.page_size,
         server_args=server_args,
         use_mla=use_mla,
+        host_size=kv_host_size,
     )
     swa_host_pool = build_kv_host_pool(
         kv_pool=swa_kv_pool,
         page_size=params.page_size,
         server_args=server_args,
         use_mla=use_mla,
+        host_size=swa_host_size,
     )
 
     # For SWA hybrid, the device alloc/free goes through the inner swa_attn_allocator

test/registered/unit/mem_cache/test_hybrid_pool_assembler.py

@@ -0,0 +1,31 @@
+"""Unit test for hybrid HiCache fixed-size budget splitting."""
+
+import unittest
+
+from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import _split_hicache_size
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class _Pool:
+    def __init__(self, kv_bytes):
+        self._kv_bytes = kv_bytes
+
+    def get_kv_size_bytes(self):
+        return self._kv_bytes
+
+
+class TestSplitHicacheSize(CustomTestCase):
+    def test_splits_total_budget_by_device_bytes(self):
+        # scalar and (k, v) tuple return shapes both supported
+        shares = _split_hicache_size(
+            100, (_Pool(75 * 10**9), _Pool((15 * 10**9, 10 * 10**9)))
+        )
+        self.assertEqual(shares, (75.0, 25.0))  # proportional to device KV bytes
+        self.assertEqual(sum(shares), 100)  # total budget preserved, not doubled
+
+
+if __name__ == "__main__":
+    unittest.main()