Fix --hicache-size allocating ~2x host memory on hybrid Mamba (#32915)

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

@@ -552,16 +552,22 @@ def build_hybrid_mamba_stack(
 ) -> tuple[HostPoolGroup, HybridCacheController]:
     transfer_layer_num = len(full_layer_mapping | mamba_layer_mapping)
     mamba_allocator = params.req_to_token_pool.mamba_allocator
+    kv_host_size, mamba_host_size = None, 0
+    if server_args.hicache_size > 0:
+        kv_host_size, mamba_host_size = _split_hicache_size(
+            server_args.hicache_size, (kv_pool, mamba_pool)
+        )
     kv_host_pool = build_kv_host_pool(
         kv_pool=kv_pool,
         page_size=params.page_size,
         server_args=server_args,
         use_mla=use_mla,
+        host_size=kv_host_size,
     )
     mamba_host_pool = MambaPoolHost(
         mamba_pool,
         server_args.hicache_ratio,
-        server_args.hicache_size,
+        mamba_host_size,
         allocator_type=_get_allocator_type(server_args),
         layout=server_args.hicache_mem_layout,
     )
@@ -639,22 +645,29 @@ def build_hybrid_mamba_swa_stack(
     )
     swa_attn_allocator = params.token_to_kv_pool_allocator.swa_attn_allocator
     mamba_allocator = params.req_to_token_pool.mamba_allocator
+    kv_host_size, swa_host_size, mamba_host_size = None, None, 0
+    if server_args.hicache_size > 0:
+        kv_host_size, swa_host_size, mamba_host_size = _split_hicache_size(
+            server_args.hicache_size, (full_kv_pool, swa_kv_pool, mamba_pool)
+        )
     kv_host_pool = build_kv_host_pool(
         kv_pool=full_kv_pool,
         page_size=page_size,
         server_args=server_args,
         use_mla=False,
+        host_size=kv_host_size,
     )
     swa_host_pool = build_kv_host_pool(
         kv_pool=swa_kv_pool,
         page_size=page_size,
         server_args=server_args,
         use_mla=False,
+        host_size=swa_host_size,
     )
     mamba_host_pool = MambaPoolHost(
         mamba_pool,
         server_args.hicache_ratio,
-        server_args.hicache_size,
+        mamba_host_size,
         allocator_type=server_args.hicache_storage_backend,
         layout=server_args.hicache_mem_layout,
     )

python/sglang/srt/mem_cache/memory_pool.py

@@ -1111,6 +1111,9 @@ class MambaPool:
             subdims_per_tensor += [subdims] * self.num_mamba_layers
         return subdims_per_tensor
 
+    def get_kv_size_bytes(self):
+        return self.mamba_cache.mem_usage_bytes()
+
 
 class HybridReqToTokenPool(ReqToTokenPool):
     """A memory pool that maps a request to its token locations."""