fix

Author: hiworldwzj

lightllm/server/embed_cache/copy_to_cache.py

@@ -0,0 +1,71 @@
+import torch
+
+import triton
+import triton.language as tl
+from typing import Optional
+
+
+@triton.jit
+def _offload_embed_tensor_to_cache(
+    embed_tensor_ptr,
+    gpu_stride0,
+    gpu_stride1,
+    gpu_stride2,
+    cache_tensor_ptr,
+    cpu_stride0,
+    cpu_stride1,
+    cpu_stride2,
+    start_index_in_cache,
+    layer_num,
+    hidden_size,
+    BLOCK: tl.constexpr,
+):
+    token_index = tl.program_id(0).to(tl.int64)
+    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
+            gpu_data = tl.load(
+                embed_tensor_ptr + token_index * gpu_stride0 + layer_index * gpu_stride1 + off * gpu_stride2, mask=mask
+            )
+            tl.store(
+                cache_tensor_ptr + dest_index * cpu_stride0 + layer_index * cpu_stride1 + off * cpu_stride2,
+                gpu_data,
+                mask=mask,
+            )
+
+    return
+
+
+@torch.no_grad()
+def offload_embed_tensor_to_cache(
+    embed_tensor: torch.Tensor,
+    cache_tensor: torch.Tensor,
+    start_index_in_cache: int,
+):
+    if len(embed_tensor.shape) == 2:
+        embed_tensor = embed_tensor.reshape(embed_tensor.shape[0], 1, embed_tensor.shape[1])
+
+    token_num = embed_tensor.shape[0]
+    grid = (token_num,)
+
+    _offload_embed_tensor_to_cache[grid](
+        embed_tensor_ptr=embed_tensor,
+        gpu_stride0=embed_tensor.stride(0),
+        gpu_stride1=embed_tensor.stride(1),
+        gpu_stride2=embed_tensor.stride(2),
+        cache_tensor_ptr=cache_tensor,
+        cpu_stride0=cache_tensor.stride(0),
+        cpu_stride1=cache_tensor.stride(1),
+        cpu_stride2=cache_tensor.stride(2),
+        start_index_in_cache=start_index_in_cache,
+        layer_num=embed_tensor.shape[1],
+        hidden_size=embed_tensor.shape[2],
+        BLOCK=256,
+        num_warps=4,
+        num_stages=1,
+    )
+    return

lightllm/server/embed_cache/embed_cache_client.py

@@ -20,19 +20,19 @@ class CpuEmbedCacheClient(object):
     This class is responsible for handling cpu kv cache meta data.
     """
 
-    def __init__(self, only_create_meta_data: bool, init_shm_data: bool):
+    def __init__(self, create_meta_data: bool, init_shm_data: bool):
         self.args = get_env_start_args()
         # to do here need calcu from from settings.
         self.embed_cache_tensor_meta = calcu_embed_cache_meta()
         self.token_num: int = self.embed_cache_tensor_meta.token_num
 
-        if not only_create_meta_data:
+        if create_meta_data:
             self.token_index_manager = MemoryManager(total_size=self.token_num)
         else:
             if init_shm_data:
                 self._create_shm_embed_kv_cache()
             else:
-                self._create_shm_embed_kv_cache()
+                self._attach_shm_cpu_embed_cache()
         return
 
     def alloc_indexes(self, token_num: int) -> Optional["MemoryBlock"]:
@@ -42,6 +42,15 @@ def release_indexes(self, block: "MemoryBlock"):
         self.token_index_manager.release(block)
         return
 
+    def copy_to_cache(self, embed_tensor: torch.Tensor, start_index_in_cache: int):
+        from .copy_to_cache import offload_embed_tensor_to_cache
+
+        offload_embed_tensor_to_cache(
+            embed_tensor=embed_tensor,
+            cache_tensor=self.cpu_embed_cache_tensor,
+            start_index_in_cache=start_index_in_cache,
+        )
+
     def _create_shm_embed_kv_cache(self):
         shm_ptr = create_shm_embed_cache_ptr()
         numpy_array = np.frombuffer(

lightllm/server/embed_cache/impl/naive_memory_cache.py

@@ -42,7 +42,7 @@ def __init__(self, args) -> None:
         self.token_id_range_start = 0
         self.token_id_range_end = 0
         self.use_config_server = self.args.config_server_host and self.args.config_server_port
-        self.cpu_embed_cache_client = CpuEmbedCacheClient(only_create_meta_data=True, init_shm_data=False)
+        self.cpu_embed_cache_client = CpuEmbedCacheClient(create_meta_data=True, init_shm_data=False)
 
     def _check_and_set_new_id_range(self, alloced_token_num):
         need_update_range = self.token_id_range_start + alloced_token_num >= self.token_id_range_end

lightllm/server/visualserver/model_infer/model_rpc.py

@@ -19,19 +19,11 @@
 from lightllm.models.qwen2_5_vl.qwen2_5_visual import Qwen2_5_VisionTransformerPretrainedModel
 from lightllm.models.qwen3_vl.qwen3_visual import Qwen3VisionTransformerPretrainedModel
 from lightllm.models.tarsier2.tarsier2_visual import TarsierVisionTransformerPretrainedModel
-from lightllm.server.embed_cache.utils import (
-    tensor2bytes,
-    read_shm,
-    create_shm,
-    create_shm_and_dump,
-    get_shm_name_data,
-    get_shm_name_embed,
-)
 from lightllm.utils.infer_utils import set_random_seed
-from lightllm.utils.infer_utils import calculate_time, mark_start, mark_end
 from lightllm.utils.dist_utils import init_vision_distributed_env
 from lightllm.utils.graceful_utils import graceful_registry
 from lightllm.utils.envs_utils import get_env_start_args
+from lightllm.server.embed_cache.embed_cache_client import CpuEmbedCacheClient
 
 
 class VisualModelRpcServer(rpyc.Service):
@@ -91,6 +83,7 @@ def exposed_init_model(self, kvargs):
 
             self.model.load_model(weight_dir)
             self.model = self.model.cuda()
+            self.cpu_embed_cache_client = CpuEmbedCacheClient(create_meta_data=False, init_shm_data=True)
         except Exception as e:
             print("#" * 16)
             print("load model error:", str(e), e, type(e))
@@ -111,7 +104,7 @@ def forward(self, images: List[ImageItem]):
     def exposed_encode(self, images: List[ImageItem]):
         images = obtain(images)
         all_img_embeds, uuids, valid_ids = self.forward(images)
-        all_img_embeds = all_img_embeds.to(torch.device("cpu"))
+        all_img_embeds = all_img_embeds.to(torch.device("cuda"))
 
         if self.tp_rank_id == 0:
             ready_flags = obtain(self.cache_client.root.get_items_embed(uuids))
@@ -121,9 +114,13 @@ def exposed_encode(self, images: List[ImageItem]):
                     continue
                 uid = uuids[i]
                 start, end = valid_ids[i]
-                create_shm_and_dump(get_shm_name_embed(uid), all_img_embeds[start:end])
+                image = images[i]
+                self.cpu_embed_cache_client.copy_to_cache(
+                    embed_tensor=all_img_embeds[start:end], start_index_in_cache=image.start_index_in_embed_cache
+                )
                 ids_to_set.append(uid)
             if ids_to_set:
+                torch.cuda.current_stream().synchronize()
                 self.cache_client.root.set_items_embed(ids_to_set)
         return