update & fix format

Author: jinbiaoyu

lightllm/common/basemodel/basemodel.py

@@ -179,28 +179,16 @@ def _init_mem_manager(self):
         )
 
         if self.enable_hiradix_cache:
-            from lightllm.common.radixmem_buffer import RadixMemoryBuffer, init_shared_data, get_shared_data, MemPropties
-            from lightllm.common.radixmem_manager import RadixBufferManager
+            from lightllm.common.radixmem_buffer import get_shared_data, MemPropties
+            from lightllm.common.radixmem_manager import build_radix_manager
             mem_propties = MemPropties(
                 self.hiradix_cache_token_num,
                 dtype=self.data_type,
                 head_num=self.config["num_attention_heads"] // self.tp_world_size_,
                 head_dim=self.config["n_embed"] // self.config["num_attention_heads"],
                 layer_num=self.config["n_layer"]
             )
-            init_shared_data(
-                mem_propties=mem_propties,
-                device="cpu" if not self.hiradix_cache_gpu else "cuda"
-            )
-            radix_mem_buffer = RadixMemoryBuffer(
-                mem_propties,
-                shared_data=get_shared_data(),
-                lock=self.radix_lock,
-                device="cpu" if not self.hiradix_cache_gpu else "cuda"
-            )
-            self.radix_manager = RadixBufferManager(radix_buffer=radix_mem_buffer, 
-                                            radix_mem_data=get_shared_data(),
-                                            lock=self.radix_lock)
+            self.radix_manager = build_radix_manager(mem_propties, self.hiradix_cache_gpu, self.radix_lock)
             self.mem_propties = mem_propties
             self.shared_mem_data = get_shared_data()
         return

lightllm/common/radixmem_buffer.py

@@ -137,14 +137,16 @@ def free_req_index(self, req_id: int):
                 return
             index = self.req_mem_index[req_id]
             self._free(index)
-            logger.info(f"Freed memory index for request {req_id} size {len(index)}, left size {self.can_use_mem_size.get_value()}")
+            logger.info(f"Freed memory index for request {req_id} size {len(index)}, "
+                        f"left size {self.can_use_mem_size.get_value()}")
             del self.req_mem_index[req_id]
 
     def alloc(self, need_size) -> torch.Tensor:
         with self.lock:
             if need_size > self.mark_end.get_value() - self.mark_start.get_value():
                 logger.error(
-                    f"warn no enough cache need_size {need_size} left_size {self.can_use_mem_size.get_value()}"
+                    f"warn no enough cache need_size {need_size} "
+                    f"left_size {self.can_use_mem_size.get_value()}"
                 )
                 raise RuntimeError(f"Not enough memory to allocate {need_size} tokens.")
 
@@ -160,7 +162,8 @@ def set_req_mem_index(self, req_id: int, index: List[int]):
         """Set the memory index for a specific request ID."""
         with self.lock:
             if req_id in self.req_mem_index:
-                logger.info(f"Request ID {req_id} already exists. Overwriting index {self.req_mem_index[req_id]} with {index}.")
+                logger.info(f"Request ID {req_id} already exists. "
+                            f"Overwriting index {self.req_mem_index[req_id]} with {index}.")
             self.req_mem_index[req_id] = index
             logger.info(f"radix mem buffer insert req {req_id}, current disk work num {self._get_current_work_num()}")
 

lightllm/common/radixmem_manager.py

@@ -6,6 +6,7 @@
 import torch.multiprocessing as mp
 from collections import OrderedDict
 
+from .radixmem_buffer import MemPropties, init_shared_data, get_shared_data
 from .radixmem_buffer import SharedRadixMemoryData, RadixMemoryBuffer
 
 from lightllm.utils.log_utils import init_logger
@@ -116,4 +117,29 @@ def query_cache(self, tokens: List[int]) -> int:
     def clear(self):
         with self.lock:
             self.radix_buffer.req_mem_index.clear()
-            self.lru_queue[:] = []
+            self.lru_queue[:] = []
+
+def build_radix_manager(mem_propties: MemPropties, 
+                        use_gpu: bool, 
+                        radix_lock) -> RadixBufferManager:
+    device = "cuda" if use_gpu else "cpu"
+
+    init_shared_data(
+        mem_propties=mem_propties,
+        device=device,
+    )
+
+    radix_mem_buffer = RadixMemoryBuffer(
+        mem_propties=mem_propties,
+        shared_data=get_shared_data(),
+        lock=radix_lock,
+        device=device,
+    )
+
+    radix_manager = RadixBufferManager(
+        radix_buffer=radix_mem_buffer,
+        radix_mem_data=get_shared_data(),
+        lock=radix_lock,
+    )
+
+    return radix_manager

lightllm/models/deepseek2/model.py

@@ -113,28 +113,16 @@ def _init_mem_manager(self):
             mem_fraction=self.mem_fraction,
         )
         if self.enable_hiradix_cache:
-            from lightllm.common.radixmem_buffer import RadixMemoryBuffer, init_shared_data, get_shared_data, MemPropties
-            from lightllm.common.radixmem_manager import RadixBufferManager
+            from lightllm.common.radixmem_buffer import get_shared_data, MemPropties
+            from lightllm.common.radixmem_manager import build_radix_manager
             mem_propties = MemPropties(
                 self.hiradix_cache_token_num,
                 dtype=self.data_type,
                 head_num=1,
                 head_dim=self.config["kv_lora_rank"] + self.config["qk_rope_head_dim"],
                 layer_num=self.config["num_hidden_layers"] + added_mtp_layer_num,
             )
-            init_shared_data(
-                mem_propties=mem_propties,
-                device="cpu" if not self.hiradix_cache_gpu else "cuda"
-            )
-            radix_mem_buffer = RadixMemoryBuffer(
-                mem_propties,
-                shared_data=get_shared_data(),
-                lock=self.radix_lock,
-                device="cpu" if not self.hiradix_cache_gpu else "cuda"
-            )
-            self.radix_manager = RadixBufferManager(radix_buffer=radix_mem_buffer, 
-                                            radix_mem_data=get_shared_data(),
-                                            lock=self.radix_lock)
+            self.radix_manager = build_radix_manager(mem_propties, self.hiradix_cache_gpu, self.radix_lock)
             self.mem_propties = mem_propties
             self.shared_mem_data = get_shared_data()
         return

lightllm/models/qwen2/model.py

@@ -53,28 +53,16 @@ def _init_mem_manager(self):
         )
 
         if self.enable_hiradix_cache:
-            from lightllm.common.radixmem_buffer import RadixMemoryBuffer, init_shared_data, get_shared_data, MemPropties
-            from lightllm.common.radixmem_manager import RadixBufferManager
+            from lightllm.common.radixmem_buffer import MemPropties, get_shared_data, MemPropties
+            from lightllm.common.radixmem_manager import build_radix_manager
             mem_propties = MemPropties(
                 self.hiradix_cache_token_num,
                 dtype=self.data_type,
                 head_num=2 * tp_k_head_num_,
                 head_dim=head_dim_,
                 layer_num=self.config["num_hidden_layers"],
             )
-            init_shared_data(
-                mem_propties=mem_propties,
-                device="cpu" if not self.hiradix_cache_gpu else "cuda"
-            )
-            radix_mem_buffer = RadixMemoryBuffer(
-                mem_propties,
-                shared_data=get_shared_data(),
-                lock=self.radix_lock,
-                device="cpu" if not self.hiradix_cache_gpu else "cuda"
-            )
-            self.radix_manager = RadixBufferManager(radix_buffer=radix_mem_buffer, 
-                                            radix_mem_data=get_shared_data(),
-                                            lock=self.radix_lock)
+            self.radix_manager = build_radix_manager(mem_propties, self.hiradix_cache_gpu, self.radix_lock)
             self.mem_propties = mem_propties
             self.shared_mem_data = get_shared_data()
         return

lightllm/server/router/model_infer/mode_backend/base_backend.py

@@ -7,7 +7,6 @@
 from lightllm.utils.log_utils import init_logger
 from lightllm.models import get_model
 from lightllm.server.router.dynamic_prompt.radix_cache import RadixCache
-from lightllm.server.router.model_infer.infer_batch import InferReq
 from lightllm.server.router.dynamic_prompt.hiradix.hiradix_cache import HiRadixCache
 from lightllm.server.router.model_infer.infer_batch import InferReq, InferSamplingParams
 from lightllm.server.router.token_load import TokenLoad