fix

Author: shihaobai

lightllm/server/embed_cache/copy_to_cache.py

@@ -24,7 +24,6 @@ def _offload_embed_tensor_to_cache(
     dest_index = (start_index_in_cache + token_index).to(tl.int64)
 
     for layer_index in range(layer_num):
-        layer_index = layer_index.to(tl.int64)
         for block_index in range(tl.cdiv(hidden_size, BLOCK)):
             off = block_index * BLOCK + tl.arange(0, BLOCK)
             mask = off < hidden_size

lightllm/server/embed_cache/impl/naive_memory_cache.py

@@ -28,6 +28,9 @@ class Record(object):
     token_num: int
     mem_block: MemoryBlock
 
+    def __hash__(self) -> int:
+        return self.id
+
 
 class InMemoryCache:
     def __init__(self, args) -> None:
@@ -162,7 +165,7 @@ def alloc(self, md5sum_list: list[str], token_num_list: list[int]) -> Optional[l
                 free_min_count=new_needed - (self.capacity - self.occupied), new_md5_dict=new_md5_dict
             )
             if len(alloc_md5_dict) == len(new_md5_dict):
-                for md5sum, mem_block in alloc_md5_dict:
+                for md5sum, mem_block in alloc_md5_dict.items():
                     token_num = new_md5_dict[md5sum]
                     uid_int = uuid.uuid1().int
                     self._check_and_set_new_id_range(token_num)

lightllm/server/router/model_infer/infer_batch.py

@@ -43,7 +43,6 @@ def register(
         radix_cache: RadixCache,
         shm_req_manager: ShmReqManager,
         vocab_size: int,
-        cpu_embed_cache_client: Optional[CpuEmbedCacheClient] = None,
     ):
         self.args = get_env_start_args()
         from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend
@@ -58,7 +57,10 @@ def register(
         self.infer_req_ids = []
 
         self.vocab_size = vocab_size
-        self.cpu_embed_cache_client = cpu_embed_cache_client
+        return
+
+    def init_cpu_embed_cache_client(self):
+        self.cpu_embed_cache_client = CpuEmbedCacheClient(create_meta_data=False, init_shm_data=False)
         return
 
     def get_overlap_stream(self) -> torch.cuda.Stream:

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

@@ -138,6 +138,9 @@ def init_model(self, kvargs):
             self.multi_level_cache_module = MultiLevelKvCacheModule(self)
             wait_events.append(self.multi_level_cache_module)
 
+        if self.args.enable_multimodal:
+            g_infer_context.init_cpu_embed_cache_client()
+
         model_cfg, _ = PretrainedConfig.get_config_dict(self.weight_dir)
 
         model_kvargs = {
@@ -187,9 +190,6 @@ def init_model(self, kvargs):
             radix_cache=self.radix_cache,
             shm_req_manager=self.shm_req_manager,
             vocab_size=self.model.vocab_size,
-            cpu_embed_cache_client=CpuEmbedCacheClient(create_meta_data=False, init_shm_data=False)
-            if self.args.enable_multimodal
-            else None,
         )
 
         # 初始化 dp 模式使用的通信 tensor, 对于非dp模式，不会使用到