add cpu embed to llm

Author: hiworldwzj

lightllm/server/router/model_infer/infer_batch.py

@@ -18,6 +18,7 @@
 from lightllm.utils.custom_kernel_utis import custom_cat
 from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.server.pd_io_struct import NIXLDecodeNodeInfo
+from lightllm.server.embed_cache.embed_cache_client import CpuEmbedCacheClient
 
 logger = init_logger(__name__)
 
@@ -30,12 +31,19 @@ class InferenceContext:
     requests_mapping: Dict[int, "InferReq"] = None
     infer_req_ids = None
     vocab_size = None
+    cpu_embed_cache_client: Optional[CpuEmbedCacheClient] = None
 
     overlap_stream: torch.cuda.Stream = None  # 一些情况下推理进程进行异步折叠操作的异步流对象。
     cpu_kv_cache_stream: torch.cuda.Stream = None  # 用 cpu kv cache 操作的 stream
 
     def register(
-        self, backend, req_manager: ReqManager, radix_cache: RadixCache, shm_req_manager: ShmReqManager, vocab_size: int
+        self,
+        backend,
+        req_manager: ReqManager,
+        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
@@ -50,6 +58,7 @@ def register(
         self.infer_req_ids = []
 
         self.vocab_size = vocab_size
+        self.cpu_embed_cache_client = cpu_embed_cache_client
         return
 
     def get_overlap_stream(self) -> torch.cuda.Stream:

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

@@ -39,6 +39,7 @@
 from lightllm.common.basemodel.triton_kernel.gather_token_id import scatter_token
 from lightllm.server.pd_io_struct import NIXLChunckedTransTaskRet
 from .multi_level_kv_cache import MultiLevelKvCacheModule
+from lightllm.server.embed_cache.embed_cache_client import CpuEmbedCacheClient
 
 
 class ModeBackend:
@@ -179,12 +180,16 @@ def init_model(self, kvargs):
             self.preload_prompt_cache_kv_buffer(model_cfg)
 
         self.logger.info(f"loaded model class {self.model.__class__}")
+
         g_infer_context.register(
             backend=self,
             req_manager=self.model.req_manager,
             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模式，不会使用到