[update] add cmm_size & precompute

Author: LittleMouse

backend/llm_client_backend.py

@@ -77,8 +77,8 @@ async def _get_client(self, request):
                 await asyncio.sleep(retry_interval)
                 await asyncio.wait_for(self._pool_lock.acquire(), timeout=timeout - (time.time() - start_time))
 
-            # if "memory_required" in self.config:
-            #     await self.memory_checker.check_memory(self.config["memory_required"])
+            if "memory_required" in self.config:
+                await self.memory_checker.check_memory(self.config["memory_required"])
 
             self.logger.debug("Creating new LLM client")
             client = LLMClient(

services/model_list.py

@@ -56,6 +56,13 @@ async def get_model_list(self, required_mem: int) -> None:
                             "object": f"{model_type}.setup",
                             "system_prompt": "You are a helpful assistant."
                         })
+
+                        mode_param = model_data.get("mode_param", {})
+                        precompute_len = None
+                        cmm_size = None
+                        if isinstance(mode_param, dict):
+                            precompute_len = mode_param.get("precompute_len")
+                            cmm_size = mode_param.get("cmm_size")
                         if '-1.5B-' in mode:
                             new_entry['memory_required'] = 1782579
                             new_entry['pool_size'] = 1
@@ -68,6 +75,8 @@ async def get_model_list(self, required_mem: int) -> None:
                         else:
                             new_entry['memory_required'] = 1363148
                             new_entry['pool_size'] = 2
+                        if cmm_size is not None:
+                            new_entry['memory_required'] = cmm_size
 
                         if '-p256-' in mode:
                             new_entry['max_context_length'] = 256
@@ -81,6 +90,9 @@ async def get_model_list(self, required_mem: int) -> None:
                         if '-ctx-' in mode:
                             new_entry['max_context_length'] = 1024
 
+                        if precompute_len is not None:
+                            new_entry['max_context_length'] = precompute_len
+
                     elif model_type == 'tts':
                         if 'melotts' in mode.lower():
                             obj = 'melotts.setup'